def _replace_tensors_with_constant_values(saved_model_dir, base_temp_dir,
*tensor_bindings):
"""Replaces specified `Tensor`s with constant values.
Constants are accepted as Python values; these are automatically
wrapped in `tf.constant()`.
This method creates its own temp dir, and is therefore idempotent
since any retry will use a different temp dir.
Args:
saved_model_dir: A SavedModel directory providing a transform
graph. The MetaGraphDef and signature are selected from the
SavedModel using keys defined in `../constants.py` ('transform'
and 'transform_signature', respectively).
base_temp_dir: Base temp dir for storage of new model.
*tensor_bindings: An iterable of `_TensorBinding`s.
Returns:
The directory name containing the updated SavedModel.
Raises:
RuntimeError: if there is no default graph available to which to
apply the transform.
"""
with tf.compat.v1.Graph().as_default() as graph:
tensor_replacement_map = {}
for tensor_binding in tensor_bindings:
# TODO(b/34792459): Make this an assertion and remove nested code once TFT
# doesn't allow missing tensor bindings (once combiner defaults are used).
if not isinstance(tensor_binding, _TensorBinding):
tf.compat.v1.logging.error(
'Encountered an empty tensor value binding, '
'is the analysis dataset empty? Tensor bindings: %s',
tensor_bindings)
assert isinstance(tensor_binding,
beam.pvalue.EmptySideInput), tensor_binding
beam.metrics.Metrics.counter(beam_common.METRICS_NAMESPACE,
'empty_tensor_bindings').inc()
continue
replacement_tensor = tf.constant(tensor_binding.value)
if tensor_binding.is_asset_filepath:
graph.add_to_collection(tf.compat.v1.GraphKeys.ASSET_FILEPATHS,
replacement_tensor)
tensor_replacement_map[
tensor_binding.tensor_name] = replacement_tensor
with tf.compat.v1.Session(graph=graph) as session:
temp_dir = beam_common.get_unique_temp_path(base_temp_dir)
input_tensors, output_tensors = (
saved_transform_io.partially_apply_saved_transform_internal(
saved_model_dir, {}, tensor_replacement_map))
session.run(tf.compat.v1.global_variables_initializer())
saved_transform_io.write_saved_transform_from_session(
session, input_tensors, output_tensors, temp_dir)
return temp_dir
def write_metadata_output(metadata):
output_path = self._path
if self._write_to_unique_subdirectory:
output_path = common.get_unique_temp_path(self._path)
metadata_io.write_metadata(metadata, output_path)
if asset_map:
with tf.io.gfile.GFile(
os.path.join(
output_path,
output_wrapper.TFTransformOutput.ASSET_MAP),
'w') as f:
f.write(json.dumps(asset_map))
return output_path
def _replace_tensors_with_constant_values(saved_model_dir, base_temp_dir,
*tensor_bindings):
"""Replaces specified `Tensor`s with constant values.
Constants are accepted as Python values; these are automatically
wrapped in `tf.constant()`.
This method creates its own temp dir, and is therefore idempotent
since any retry will use a different temp dir.
Args:
saved_model_dir: A SavedModel directory providing a transform
graph. The MetaGraphDef and signature are selected from the
SavedModel using keys defined in `../constants.py` ('transform'
and 'transform_signature', respectively).
base_temp_dir: Base temp dir for storage of new model.
*tensor_bindings: An iterable of `_TensorBinding`s.
Returns:
The directory name containing the updated SavedModel.
Raises:
RuntimeError: if there is no default graph available to which to
apply the transform.
"""
with tf.compat.v1.Graph().as_default() as graph:
tensor_replacement_map = {}
for tensor_binding in tensor_bindings:
assert isinstance(tensor_binding, _TensorBinding), tensor_binding
replacement_tensor = tf.constant(tensor_binding.value)
if tensor_binding.is_asset_filepath:
graph.add_to_collection(tf.compat.v1.GraphKeys.ASSET_FILEPATHS,
replacement_tensor)
tensor_replacement_map[
tensor_binding.tensor_name] = replacement_tensor
with tf.compat.v1.Session(graph=graph) as session:
temp_dir = beam_common.get_unique_temp_path(base_temp_dir)
input_tensors, output_tensors = (
saved_transform_io.partially_apply_saved_transform_internal(
saved_model_dir, {}, tensor_replacement_map))
session.run(tf.compat.v1.global_variables_initializer())
saved_transform_io.write_saved_transform_from_session(
session, input_tensors, output_tensors, temp_dir)
return temp_dir
def _create_saved_model_impl(inputs, operation, extra_args):
"""Create a SavedModel from a TF Graph."""
unbound_saved_model_dir = common.get_unique_temp_path(
extra_args.base_temp_dir)
with extra_args.graph.as_default():
with tf.Session(graph=extra_args.graph) as session:
table_initializers_ref = tf.get_collection_ref(
tf.GraphKeys.TABLE_INITIALIZERS)
original_table_initializers = list(table_initializers_ref)
del table_initializers_ref[:]
table_initializers_ref.extend(operation.table_initializers)
# Initialize all variables so they can be saved.
session.run(tf.global_variables_initializer())
saved_transform_io.write_saved_transform_from_session(
session, extra_args.input_signature, operation.output_signature,
unbound_saved_model_dir)
del table_initializers_ref[:]
table_initializers_ref.extend(original_table_initializers)
return inputs | operation.label >> _BindTensors(
extra_args.base_temp_dir, unbound_saved_model_dir, extra_args.pipeline)
def expand(self, inputs):
unbound_saved_model_dir = beam_common.get_unique_temp_path(
self._base_temp_dir)
with self._graph.as_default():
with tf.compat.v1.Session(graph=self._graph) as session:
table_initializers_ref = tf.compat.v1.get_collection_ref(
tf.compat.v1.GraphKeys.TABLE_INITIALIZERS)
original_table_initializers = list(table_initializers_ref)
del table_initializers_ref[:]
table_initializers_ref.extend(self._table_initializers)
# Initialize all variables so they can be saved.
session.run(tf.compat.v1.global_variables_initializer())
saved_transform_io.write_saved_transform_from_session(
session, self._input_signature, self._output_signature,
unbound_saved_model_dir)
del table_initializers_ref[:]
table_initializers_ref.extend(original_table_initializers)
return (inputs
| 'BindTensors' >> _BindTensors(self._base_temp_dir,
unbound_saved_model_dir)
| 'Count' >> beam_common.IncrementCounter('saved_models_created'))
def _copy_tree_to_unique_temp_dir(source, base_temp_dir_path):
"""Copies from source to a unique sub directory under base_temp_dir_path."""
destination = common.get_unique_temp_path(base_temp_dir_path)
_copy_tree(source, destination)
return destination
def write_metadata_output(metadata):
output_path = self._path
if self._write_to_unique_subdirectory:
output_path = common.get_unique_temp_path(self._path)
metadata_io.write_metadata(metadata, output_path)
return output_path
