def _get_temporary_analyzer_output(
temp_dir: str,
tensor_info: TensorInfo,
name: Optional[str] = None) -> TemporaryAnalyzerOutputWrapper:
"""Create a temporary graph tensor using attributes in `tensor_info`.
Args:
temp_dir: Path to a directory to write out any temporary asset files to.
tensor_info: A `TensorInfo` object containing attributes to create the graph
tensor.
name: A string (or None). The created graph tensor uses this name.
Returns:
A named tuple `TemporaryAnalyzerOutputWrapper` with:
asset: If the graph tensor represents a path to an asset file, a
`tf.saved_model.Asset` object for tracking. Else, None.
graph_tensor: The graph tensor
"""
asset = None
with tf.name_scope('temporary_analyzer_output'):
is_asset_filepath = tensor_info.temporary_asset_value is not None
if is_asset_filepath:
# Placeholders cannot be used for assets, if this graph will be serialized
# to a SavedModel, as they will be initialized with the init op. If a
# `temp_dir` is provided, it is assumed that this graph will be
# serialized and a temporary asset file is written out . Else, a
# placeholder is returned.
# TODO(b/164921571) Support temporary files in tfrecord format.
# TODO(b/149997088): Reduce number of temporary files written out.
if temp_dir:
with tf.init_scope():
temporary_asset_filepath = os.path.join(temp_dir, uuid.uuid4().hex)
with tf.io.gfile.GFile(temporary_asset_filepath, 'w') as f:
f.write(tensor_info.temporary_asset_value)
# Wrap asset files using `tf.saved_model.Asset` to ensure that
# `SavedModel`s exported are hermetic.
asset = common_types.Asset(temporary_asset_filepath)
graph_tensor = tf.constant(
temporary_asset_filepath,
dtype=tensor_info.dtype,
shape=tensor_info.shape,
name=name)
else:
graph_tensor = tf.raw_ops.Placeholder(
dtype=tensor_info.dtype, shape=tensor_info.shape, name=name)
else:
# Using a placeholder with no default value causes tracing to fail if
# there is any control flow dependent on a child tensor of this
# placeholder. Hence, provide a temporary default value for it.
# If dtype is string, we want a tensor that contains '0's instead of b'[]
# to allow string to numeric conversion ops to trace successfully.
temporary_dtype = (
tf.int64 if tensor_info.dtype == tf.string else tensor_info.dtype)
temporary_tensor = tf2_utils.supply_missing_tensor(
1, tf.TensorShape(tensor_info.shape), temporary_dtype)
if tensor_info.dtype == tf.string:
temporary_tensor = tf.strings.as_string(temporary_tensor)
graph_tensor = tf.raw_ops.PlaceholderWithDefault(
input=temporary_tensor, shape=tensor_info.shape, name=name)
return TemporaryAnalyzerOutputWrapper(asset, graph_tensor)
def write_v2_saved_model(tf_function: function.Function, name: str,
saved_model_dir: str) -> function.ConcreteFunction:
"""Writes `tf_function` under attr `name` to `saved_model_dir`."""
module = tf.Module()
resource_tracker = tracking.ResourceTracker()
object_tracker = annotators.ObjectTracker()
created_variables = []
def _variable_creator(next_creator, **kwargs):
var = next_creator(**kwargs)
created_variables.append(var)
return var
# TODO(b/164921571): Handle generic Trackable objects.
# Trace `tf_function` to gather any resources in it using the
# resource_tracker. These are then assigned to `module.resources` and tracked
# before exporting to SavedModel.
with tracking.resource_tracker_scope(resource_tracker), \
annotators.object_tracker_scope(object_tracker), \
tf.variable_creator_scope(_variable_creator):
concrete_fn = tf_function.get_concrete_function()
# Prior to 2020/10/08, saving a tf.function with a concrete function signature
# would ensure that the function was not re-traced in a round-trip to a
# SavedModel. Since this is no longer the case, we save the concrete function
# directly.
if tf.compat.forward_compatible(2020, 10, 8):
pruned_function = optimize_concrete_function(concrete_fn)
module.pruned_variables = pruned_function.variables
setattr(module, name, pruned_function)
else:
setattr(module, name, tf_function)
# Any variables created need to be explicitly tracked.
module.created_variables = created_variables
# Resources need to be explicitly tracked.
module.resources = resource_tracker.resources
module.trackable_objects = object_tracker.trackable_objects
# TODO(b/158011374) - Stop explicitly tracking initializers. Tracking the
# table should be sufficient.
initializers = []
for resource in module.resources:
if isinstance(resource, lookup_ops.InitializableLookupTableBase):
initializers.append(resource._initializer) # pylint: disable=protected-access
module.initializers = initializers
module.assets = [
common_types.Asset(asset_filepath)
for asset_filepath in concrete_fn.graph.get_collection(
tf.compat.v1.GraphKeys.ASSET_FILEPATHS)
]
tf.saved_model.save(module, saved_model_dir)
return concrete_fn
def trace_and_update_module(
module: tf.Module, tf_function: function.Function, name: str,
strip_control_dependencies: bool) -> function.ConcreteFunction:
"""Traces `tf_function` and saves under attr `name` of `module`.
Args:
module: A saveable module which will contain the traced `tf_function` under
attr `name`.
tf_function: A tf.function to trace.
name: A name to same the traced `tf_function` to.
strip_control_dependencies: Boolean. If True, automatic control dependencies
will be stripped from the outputs of `tf_function`. This should almost
always be False. It is useful only if you want to use the structure of the
TF graph to perform any graph manipulations.
Returns:
The concrete function obtained from tracing `tf_function`.
"""
resource_tracker = tracking.ResourceTracker()
object_tracker = annotators.ObjectTracker()
created_variables = []
def _variable_creator(next_creator, **kwargs):
var = next_creator(**kwargs)
created_variables.append(var)
return var
# Trace `tf_function` to gather any resources in it using the
# resource_tracker. These are then assigned to `module.resources` and tracked
# before exporting to SavedModel.
with tracking.resource_tracker_scope(resource_tracker), \
annotators.object_tracker_scope(object_tracker), \
tf.variable_creator_scope(_variable_creator):
concrete_fn = tf_function.get_concrete_function()
# Prior to 2020/10/08, saving a tf.function with a concrete function signature
# would ensure that the function was not re-traced in a round-trip to a
# SavedModel. Since this is no longer the case, we save the concrete function
# directly.
if tf.compat.forward_compatible(2020, 10, 8):
pruned_function = optimize_concrete_function(concrete_fn,
strip_control_dependencies)
module.pruned_variables = pruned_function.variables
setattr(module, name, pruned_function)
else:
setattr(module, name, tf_function)
# Any variables created need to be explicitly tracked.
module.created_variables = created_variables
# Resources need to be explicitly tracked.
module.resources = resource_tracker.resources
module.trackable_objects = object_tracker.trackable_objects
# TODO(b/158011374) - Stop explicitly tracking initializers. Tracking the
# table should be sufficient.
initializers = []
for resource in module.resources:
if isinstance(resource, lookup_ops.InitializableLookupTableBase):
initializers.append(resource._initializer) # pylint: disable=protected-access
module.initializers = initializers
module.assets = [
common_types.Asset(asset_filepath) for asset_filepath in
concrete_fn.graph.get_collection(tf.compat.v1.GraphKeys.ASSET_FILEPATHS)
]
return concrete_fn
def trace_and_write_v2_saved_model(saved_model_dir, preprocessing_fn,
input_signature, base_temp_dir,
tensor_replacement_map,
output_keys_to_name_map):
"""Writes out a SavedModelV2 with preprocessing_fn traced using tf.function.
The SavedModel written contains a method called `transform_fn` that
represents the traced `preprocessing_fn`. Additionally, if this is the final
SavedModel being written out, it will contain a method called `metadata_fn`
that provides deferred schema annotations.
Args:
saved_model_dir: Path to write SavedModel to.
preprocessing_fn: A user defined python function to be traced.
input_signature: TypeSpecs describing the inputs to the `preprocessing_fn`.
base_temp_dir: Base path to write temporary artifacts to.
tensor_replacement_map: A map from placeholder tensor names to their
evaluated replacement tensors.
output_keys_to_name_map: A map from output dictionary keys to the names of
the tensors that they represent.
Returns:
A tuple containing a pair of `tf.ConcreteFunction`s:
1. The traced preprocessing_fn.
2. A metadata_fn that returns a dictionary containing the deferred
annotations added to the graph when invoked with any valid input.
"""
module = tf.Module()
transform_fn = get_traced_transform_fn(
preprocessing_fn,
input_signature,
base_temp_dir,
tensor_replacement_map=tensor_replacement_map,
output_keys_to_name_map=output_keys_to_name_map)
metadata_fn = None
resource_tracker = tracking.ResourceTracker()
created_variables = []
def _variable_creator(next_creator, **kwargs):
var = next_creator(**kwargs)
created_variables.append(var)
return var
# TODO(b/164921571): Handle generic Trackable objects.
# Trace the `transform_fn` and `metadata_fn` to gather any resources in it
# using the resource_tracker. These are then assigned to `module.resources`
# and tracked before exporting to SavedModel.
with tracking.resource_tracker_scope(
resource_tracker), tf.variable_creator_scope(_variable_creator):
concrete_transform_fn = transform_fn.get_concrete_function()
concrete_metadata_fn = None
# If the `TENSOR_REPLACEMENTS` graph collection is empty, all TFT analyzers
# in the `preprocessing_fn` have already been evaluated.
if not concrete_transform_fn.graph.get_collection(
analyzer_nodes.TENSOR_REPLACEMENTS):
metadata_fn = schema_inference.get_traced_metadata_fn(
tensor_replacement_map,
preprocessing_fn,
input_signature,
base_temp_dir,
evaluate_schema_overrides=True)
concrete_metadata_fn = metadata_fn.get_concrete_function()
# Save ConcreteFunction when possible since the above workaround won't work if
# the tf.function is retraced.
if tf.compat.forward_compatible(2020, 10, 8):
module.transform_fn = concrete_transform_fn
module.metadata_fn = concrete_metadata_fn
else:
module.transform_fn = transform_fn
module.metadata_fn = metadata_fn
# Any variables created need to be explicitly tracked.
module.created_variables = created_variables
# Resources need to be explicitly tracked.
module.resources = resource_tracker.resources
# TODO(b/158011374) - Stop explicitly tracking initializers. Tracking the
# table should be sufficient.
initializers = []
for resource in module.resources:
if isinstance(resource, lookup_ops.InitializableLookupTableBase):
initializers.append(resource._initializer) # pylint: disable=protected-access
module.initializers = initializers
module.assets = [
common_types.Asset(asset_filepath)
for asset_filepath in concrete_transform_fn.graph.get_collection(
tf.compat.v1.GraphKeys.ASSET_FILEPATHS)
]
tf.saved_model.save(module, saved_model_dir)
return concrete_transform_fn, concrete_metadata_fn
def _bind_future_as_tensor_v2(
future: nodes.ValueNode,
tensor_info: TensorInfo,
name: Optional[str] = None) -> common_types.TemporaryAnalyzerOutputType:
"""Bind a future value as a tensor to a TF2 FuncGraph.
If the future is expected to write out an asset file and this method is
invoked within a `TFGraphContext` that was provided a temporary directory,
a temporary file is written out by this method.
This could write out a significant number of temporary files depending on
number of times the `preprocessing_fn` is traced and number of asset files
in each tracing.
Args:
future: Future whose result should replace the graph tensor to which its
bound.
tensor_info: A `TensorInfo` object containing attributes to create the graph
tensor.
name: (Optional) If provided, the graph tensor created uses this name.
Returns:
A graph tensor or `tf.saved_model.Asset` that this future is bound to. If
this future has already been evaluated in a previous TFT phase, it is
directly returned.
"""
graph = ops.get_default_graph()
temp_dir = TFGraphContext.get_or_create_temp_dir()
temporary_analyzer_info = _get_temporary_analyzer_output(
temp_dir, tensor_info, name)
is_asset_filepath = tensor_info.temporary_asset_value is not None
# TODO(b/149997088): Switch to using a counter instead of tensor names.
# Check if an evaluated value exists for this analyzer node.
evaluated_replacements = TFGraphContext.get_evaluated_replacements()
# evaluated_replacements is a dictionary from placeholder name to evaluated
# tensor.
# If `preprocessing_fn` was traced previously and this future was then
# evaluated in a TFT phase, the result will be present in this dictionary.
analyzer_name = temporary_analyzer_info.graph_tensor.name
if (evaluated_replacements is not None and
analyzer_name in evaluated_replacements):
replaced_result = evaluated_replacements[analyzer_name]
if is_asset_filepath:
# Wrap asset files using `tf.saved_model.Asset` to ensure that
# `SavedModel`s exported are hermetic.
with tf.init_scope():
asset = common_types.Asset(replaced_result)
graph.add_to_collection(ASSET_REPLACEMENTS,
TemporaryAnalyzerOutputWrapper(asset, None))
return asset
else:
return replaced_result
else:
graph.add_to_collection(
TENSOR_REPLACEMENTS,
TensorSink(temporary_analyzer_info.graph_tensor, future,
is_asset_filepath))
if is_asset_filepath and temporary_analyzer_info.asset:
graph.add_to_collection(ASSET_REPLACEMENTS, temporary_analyzer_info)
return temporary_analyzer_info.asset
else:
return temporary_analyzer_info.graph_tensor