I have a TensorFlow code that involves using TFTransformOutput to transform raw features read from an existing recordIO of Tf.Examples and then feeding those transformed features into a warm started TensorFlow estimator.EstimatorV2 for predictions.

The entire process is managed within a single TensorFlow graph and session context to avoid graph mismatch errors potentially due to using both tf_v1 and tf_v2 in the code.

However, I keep encountering an error related to tensors being from different graphs. Below is the code and the error message I’m facing. I would really appreciate your help.

Trimmed code:

def main(argv: Sequence[str]) -> None:
...
  tf.compat.v1.disable_v2_behavior()

  def transform_features_fn(
      tf_transform_output,
      hparams,
      examples,
      used_features,
  ):
    """Creates an input function for serving based on ExampleMetadata."""
    feature_spec = tf_transform_output.raw_feature_spec()

    parsed_features = tf.io.parse_example(examples, feature_spec)
    transformed_features = tf_transform_output.transform_raw_features(
        parsed_features, drop_unused_features=True
    )
    return transformed_features

  # Define paths and parameters
  tf_transform_output_path = "/some/path"
  checkpoint_path = "some/path/2"
  testdata_filepattern = "some/path/3"

  # Load the transform output
  tf_transform_output = trainer_util.TFTransformOutput(tf_transform_output_path)

  # Define hyperparameters
  hparams = model_hparams.create_hparams()
  # Function to get records
  def get_records(records_list):
    for f in gfile.Glob(testdata_filepattern):
      with recordio.RecordReader(f) as rr:
        for record in rr:
          ex = example_pb2.Example()
          ex.ParseFromString(record)
          records_list.append(ex)
          return  # Return after one example for simplicity

  # Get records
  records_list = []
  get_records(records_list=records_list)

  # Define the graph context
  graph = tf.Graph()
  with graph.as_default():
    with tf.compat.v1.Session(graph=graph) as sess:
      # Serialize the example
      serialized_examples_tensor = tf.constant(
          [records_list[0].SerializeToString()]
      )  # Changed to batch of one serialized example
      # Transform features within the same graph
      transformed_features = transform_features_fn(
          tf_transform_output,
          hparams,
          serialized_examples_tensor,
          used_features,
      )

      # Define input function
      def input_fn():
        dataset = tf.data.Dataset.from_tensor_slices(transformed_features)
        dataset = dataset.batch(1)
        return dataset

      # Create the estimator within the same graph
      my_model = model.create_estimator(
          tf_transform_output,
          hparams,
          warm_start_from=checkpoint_path,
          model_dir=checkpoint_path,
      )
      predictions = my_model.predict(input_fn=input_fn)
      for prediction in predictions:
        print(prediction["predictions"])

Error Message:

E0603 04:04:15.467269 3556985 app.py:668] Top-level exception: Tensor(“batch_size:0”, shape=(), dtype=int64, device=/device:CPU:0) must be from the same graph as Tensor(“TensorSliceDataset:0”, shape=(), dtype=variant) (graphs are <tensorflow.python.framework.ops.Graph object at 0x51cc1ab7b9c0> and <tensorflow.python.framework.ops.Graph object at 0x51cc24427bc0>).
…….
raise ValueError(
ValueError: Tensor(“batch_size:0”, shape=(), dtype=int64, device=/device:CPU:0) must be from the same graph as Tensor(“TensorSliceDataset:0”, shape=(), dtype=variant) (graphs are <tensorflow.python.framework.ops.Graph object at 0x51cc1ab7b9c0> and <tensorflow.python.framework.ops.Graph object at 0x51cc24427bc0>).

I tried to put all the operations in a single TensorFlow graph and session context to avoid graph mismatch errors potentially but it didn’t work.