tft.make_and_track_object

Keeps track of the object created by invoking trackable_factory_callable.

tft.make_and_track_object(
    trackable_factory_callable: Callable[[], base.Trackable],
    name: Optional[str] = None
) -> base.Trackable

This API is only for use when Transform APIs are run with TF2 behaviors enabled and tft_beam.Context.force_tf_compat_v1 is set to False.

Use this API to track TF Trackable objects created in the preprocessing_fn such as tf.hub modules, tf.data.Dataset etc. This ensures they are serialized correctly when exporting to SavedModel.

trackable_factory_callable A callable that creates and returns a Trackable object.
name (Optional) Provide a unique name to track this object with. If the Trackable object created is a Keras Layer or Model this is needed for proper tracking.

Example:

def preprocessing_fn(inputs):
  dataset = tft.make_and_track_object(
      lambda: tf.data.Dataset.from_tensor_slices([1, 2, 3]))
  with tf.init_scope():
    dataset_list = list(dataset.as_numpy_iterator())
  return {'x_0': dataset_list[0] + inputs['x']}
raw_data = [dict(x=1), dict(x=2), dict(x=3)]
feature_spec = dict(x=tf.io.FixedLenFeature([], tf.int64))
raw_data_metadata = tft.DatasetMetadata.from_feature_spec(feature_spec)
with tft_beam.Context(temp_dir=tempfile.mkdtemp(),
                      force_tf_compat_v1=False):
  transformed_dataset, transform_fn = (
      (raw_data, raw_data_metadata)
      | tft_beam.AnalyzeAndTransformDataset(preprocessing_fn))
transformed_data, transformed_metadata = transformed_dataset
transformed_data
[{'x_0': 2}, {'x_0': 3}, {'x_0': 4}]

The object returned when trackable_factory_callable is invoked. The object creation is lifted out to the eager context using tf.init_scope.