def test_build_all_signature_defs_export_outputs_required(self):
receiver_tensor = constant_op.constant(["11"])
with self.assertRaises(ValueError) as e:
export.build_all_signature_defs(receiver_tensor, None)
self.assertEqual("export_outputs must be a dict.", str(e.exception))
def test_build_all_signature_defs_serving_only(self):
receiver_tensor = {"input": array_ops.placeholder(dtypes.string)}
output_1 = constant_op.constant([1.])
export_outputs = {
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
export_output.PredictOutput(outputs=output_1),
"train": export_output.TrainOutput(loss=output_1),
}
signature_defs = export.build_all_signature_defs(
receiver_tensor, export_outputs)
expected_signature_defs = {
"serving_default": signature_def_utils.predict_signature_def(
receiver_tensor, {"output": output_1})
}
self.assertDictEqual(expected_signature_defs, signature_defs)
signature_defs = export.build_all_signature_defs(
receiver_tensor, export_outputs, serving_only=False)
expected_signature_defs.update({
"train": signature_def_utils.supervised_train_signature_def(
receiver_tensor, loss={"loss": output_1})
})
self.assertDictEqual(expected_signature_defs, signature_defs)
def test_build_all_signature_defs_serving_only(self):
receiver_tensor = {"input": array_ops.placeholder(dtypes.string)}
output_1 = constant_op.constant([1.])
export_outputs = {
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
export_output.PredictOutput(outputs=output_1),
"train":
export_output.TrainOutput(loss=output_1),
}
signature_defs = export.build_all_signature_defs(
receiver_tensor, export_outputs)
expected_signature_defs = {
"serving_default":
signature_def_utils.predict_signature_def(receiver_tensor,
{"output": output_1})
}
self.assertDictEqual(expected_signature_defs, signature_defs)
signature_defs = export.build_all_signature_defs(receiver_tensor,
export_outputs,
serving_only=False)
expected_signature_defs.update({
"train":
signature_def_utils.supervised_train_signature_def(
receiver_tensor, loss={"loss": output_1})
})
self.assertDictEqual(expected_signature_defs, signature_defs)
def test_build_all_signature_defs_export_outputs_required(self):
receiver_tensor = constant_op.constant(["11"])
with self.assertRaises(ValueError) as e:
export.build_all_signature_defs(receiver_tensor, None)
self.assertEqual("export_outputs must be a dict.", str(e.exception))
def test_build_all_signature_defs_without_receiver_alternatives(self):
receiver_tensor = array_ops.placeholder(dtypes.string)
output_1 = constant_op.constant([1.])
output_2 = constant_op.constant(["2"])
output_3 = constant_op.constant(["3"])
export_outputs = {
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
export_output.RegressionOutput(value=output_1),
"head-2": export_output.ClassificationOutput(classes=output_2),
"head-3": export_output.PredictOutput(outputs={
"some_output_3": output_3
}),
}
signature_defs = export.build_all_signature_defs(
receiver_tensor, export_outputs)
expected_signature_defs = {
"serving_default":
signature_def_utils.regression_signature_def(receiver_tensor,
output_1),
"head-2":
signature_def_utils.classification_signature_def(receiver_tensor,
output_2, None),
"head-3":
signature_def_utils.predict_signature_def({
"input": receiver_tensor
}, {"some_output_3": output_3})
}
self.assertDictEqual(expected_signature_defs, signature_defs)
def test_build_all_signature_defs_with_single_alternatives(self):
receiver_tensor = array_ops.placeholder(dtypes.string)
receiver_tensors_alternative_1 = array_ops.placeholder(dtypes.int64)
receiver_tensors_alternative_2 = array_ops.sparse_placeholder(
dtypes.float32)
# Note we are passing single Tensors as values of
# receiver_tensors_alternatives, where normally that is a dict.
# In this case a dict will be created using the default receiver tensor
# name "input".
receiver_tensors_alternatives = {"other1": receiver_tensors_alternative_1,
"other2": receiver_tensors_alternative_2}
output_1 = constant_op.constant([1.])
output_2 = constant_op.constant(["2"])
output_3 = constant_op.constant(["3"])
export_outputs = {
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
export_output.RegressionOutput(value=output_1),
"head-2": export_output.ClassificationOutput(classes=output_2),
"head-3": export_output.PredictOutput(outputs={
"some_output_3": output_3
}),
}
signature_defs = export.build_all_signature_defs(
receiver_tensor, export_outputs, receiver_tensors_alternatives)
expected_signature_defs = {
"serving_default":
signature_def_utils.regression_signature_def(
receiver_tensor,
output_1),
"head-2":
signature_def_utils.classification_signature_def(
receiver_tensor,
output_2, None),
"head-3":
signature_def_utils.predict_signature_def(
{"input": receiver_tensor},
{"some_output_3": output_3}),
"other1:head-3":
signature_def_utils.predict_signature_def(
{"input": receiver_tensors_alternative_1},
{"some_output_3": output_3}),
"other2:head-3":
signature_def_utils.predict_signature_def(
{"input": receiver_tensors_alternative_2},
{"some_output_3": output_3})
# Note that the alternatives 'other:serving_default' and 'other:head-2'
# are invalid, because regession and classification signatures must take
# a single string input. Here we verify that these invalid signatures
# are not included in the export.
}
self.assertDictEqual(expected_signature_defs, signature_defs)
def _create_signature_def_map(model, mode):
"""Create a SignatureDef map from a Keras model."""
inputs_dict = {name: x for name, x in zip(model.input_names, model.inputs)}
if model.optimizer:
targets_dict = {x.name.split(':')[0]: x
for x in model.targets if x is not None}
inputs_dict.update(targets_dict)
outputs_dict = {name: x
for name, x in zip(model.output_names, model.outputs)}
export_outputs = model_fn_lib.export_outputs_for_mode(
mode,
predictions=outputs_dict,
loss=model.total_loss if model.optimizer else None,
metrics=estimator_keras_util._convert_keras_metrics_to_estimator(model))
return export_helpers.build_all_signature_defs(
inputs_dict,
export_outputs=export_outputs,
serving_only=(mode == model_fn_lib.ModeKeys.PREDICT))
def _create_signature_def_map(model, mode):
"""Create a SignatureDef map from a Keras model."""
inputs_dict = {name: x for name, x in zip(model.input_names, model.inputs)}
if model.optimizer:
targets_dict = {
x.name.split(':')[0]: x
for x in model.targets if x is not None
}
inputs_dict.update(targets_dict)
outputs_dict = {
name: x
for name, x in zip(model.output_names, model.outputs)
}
metrics = estimator_keras_util._convert_keras_metrics_to_estimator(model)
# Add metric variables to the `LOCAL_VARIABLES` collection. Metric variables
# are by default not added to any collections. We are doing this here, so
# that metric variables get initialized.
local_vars = set(ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES))
vars_to_add = set()
if metrics is not None:
for key, value in six.iteritems(metrics):
if isinstance(value, Metric):
vars_to_add.update(value.variables)
# Convert Metric instances to (value_tensor, update_op) tuple.
metrics[key] = (value.result(), value.updates[0])
# Remove variables that are in the local variables collection already.
vars_to_add = vars_to_add.difference(local_vars)
for v in vars_to_add:
ops.add_to_collection(ops.GraphKeys.LOCAL_VARIABLES, v)
export_outputs = model_fn_lib.export_outputs_for_mode(
mode,
predictions=outputs_dict,
loss=model.total_loss if model.optimizer else None,
metrics=metrics)
return export_helpers.build_all_signature_defs(
inputs_dict,
export_outputs=export_outputs,
serving_only=(mode == model_fn_lib.ModeKeys.PREDICT))
def _create_signature_def_map(model, mode):
"""Create a SignatureDef map from a Keras model."""
inputs_dict = {name: x for name, x in zip(model.input_names, model.inputs)}
if model.optimizer:
targets_dict = {x.name.split(':')[0]: x
for x in model.targets if x is not None}
inputs_dict.update(targets_dict)
outputs_dict = {name: x
for name, x in zip(model.output_names, model.outputs)}
metrics = estimator_keras_util._convert_keras_metrics_to_estimator(model)
# Add metric variables to the `LOCAL_VARIABLES` collection. Metric variables
# are by default not added to any collections. We are doing this here, so
# that metric variables get initialized.
local_vars = set(ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES))
vars_to_add = set()
if metrics is not None:
for key, value in six.iteritems(metrics):
if isinstance(value, Metric):
vars_to_add.update(value.variables)
# Convert Metric instances to (value_tensor, update_op) tuple.
metrics[key] = (value.result(), value.updates[0])
# Remove variables that are in the local variables collection already.
vars_to_add = vars_to_add.difference(local_vars)
for v in vars_to_add:
ops.add_to_collection(ops.GraphKeys.LOCAL_VARIABLES, v)
export_outputs = model_fn_lib.export_outputs_for_mode(
mode,
predictions=outputs_dict,
loss=model.total_loss if model.optimizer else None,
metrics=metrics)
return export_helpers.build_all_signature_defs(
inputs_dict,
export_outputs=export_outputs,
serving_only=(mode == model_fn_lib.ModeKeys.PREDICT))
def export_savedmodel(self,
export_dir_base,
serving_input_receiver_fn,
assets_extra=None,
as_text=False,
checkpoint_path=None):
"""Exports inference graph as a SavedModel into given dir.
This method builds a new graph by first calling the
serving_input_receiver_fn to obtain feature `Tensor`s, and then calling
this `Estimator`'s model_fn to generate the model graph based on those
features. It restores the given checkpoint (or, lacking that, the most
recent checkpoint) into this graph in a fresh session. Finally it creates
a timestamped export directory below the given export_dir_base, and writes
a `SavedModel` into it containing a single `MetaGraphDef` saved from this
session.
The exported `MetaGraphDef` will provide one `SignatureDef` for each
element of the export_outputs dict returned from the model_fn, named using
the same keys. One of these keys is always
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, indicating which
signature will be served when a serving request does not specify one.
For each signature, the outputs are provided by the corresponding
`ExportOutput`s, and the inputs are always the input receivers provided by
the serving_input_receiver_fn.
Extra assets may be written into the SavedModel via the extra_assets
argument. This should be a dict, where each key gives a destination path
(including the filename) relative to the assets.extra directory. The
corresponding value gives the full path of the source file to be copied.
For example, the simple case of copying a single file without renaming it
is specified as `{'my_asset_file.txt': '/path/to/my_asset_file.txt'}`.
Args:
export_dir_base: A string containing a directory in which to create
timestamped subdirectories containing exported SavedModels.
serving_input_receiver_fn: A function that takes no argument and
returns a `ServingInputReceiver`.
assets_extra: A dict specifying how to populate the assets.extra directory
within the exported SavedModel, or `None` if no extra assets are needed.
as_text: whether to write the SavedModel proto in text format.
checkpoint_path: The checkpoint path to export. If `None` (the default),
the most recent checkpoint found within the model directory is chosen.
Returns:
The string path to the exported directory.
Raises:
ValueError: if no serving_input_receiver_fn is provided, no export_outputs
are provided, or no checkpoint can be found.
"""
if serving_input_receiver_fn is None:
raise ValueError('serving_input_receiver_fn must be defined.')
with ops.Graph().as_default() as g:
training.create_global_step(g)
random_seed.set_random_seed(self._config.tf_random_seed)
serving_input_receiver = serving_input_receiver_fn()
# Call the model_fn and collect the export_outputs.
estimator_spec = self._call_model_fn(
features=serving_input_receiver.features,
labels=None,
mode=model_fn_lib.ModeKeys.PREDICT)
# Build the SignatureDefs from receivers and all outputs
signature_def_map = build_all_signature_defs(
serving_input_receiver.receiver_tensors,
estimator_spec.export_outputs)
if not checkpoint_path:
# Locate the latest checkpoint
checkpoint_path = saver.latest_checkpoint(self._model_dir)
if not checkpoint_path:
raise ValueError("Couldn't find trained model at %s." %
self._model_dir)
export_dir = get_timestamped_export_dir(export_dir_base)
# TODO(soergel): Consider whether MonitoredSession makes sense here
with tf_session.Session() as session:
saver_for_restore = estimator_spec.scaffold.saver or saver.Saver(
sharded=True)
saver_for_restore.restore(session, checkpoint_path)
# TODO(b/36111876): replace legacy_init_op with main_op mechanism
# pylint: disable=protected-access
local_init_op = (
estimator_spec.scaffold.local_init_op
or monitored_session.Scaffold._default_local_init_op())
# pylint: enable=protected-access
# Perform the export
builder = saved_model_builder.SavedModelBuilder(export_dir)
builder.add_meta_graph_and_variables(
session, [tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=ops.get_collection(
ops.GraphKeys.ASSET_FILEPATHS),
legacy_init_op=local_init_op)
builder.save(as_text)
# Add the extra assets
if assets_extra:
assets_extra_path = os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes('assets.extra'))
for dest_relative, source in assets_extra.items():
dest_absolute = os.path.join(
compat.as_bytes(assets_extra_path),
compat.as_bytes(dest_relative))
dest_path = os.path.dirname(dest_absolute)
gfile.MakeDirs(dest_path)
gfile.Copy(source, dest_absolute)
return export_dir
def export_savedmodel(
self, export_dir_base, serving_input_receiver_fn,
assets_extra=None,
as_text=False,
checkpoint_path=None):
"""Exports inference graph as a SavedModel into given dir.
This method builds a new graph by first calling the
serving_input_receiver_fn to obtain feature `Tensor`s, and then calling
this `Estimator`'s model_fn to generate the model graph based on those
features. It restores the given checkpoint (or, lacking that, the most
recent checkpoint) into this graph in a fresh session. Finally it creates
a timestamped export directory below the given export_dir_base, and writes
a `SavedModel` into it containing a single `MetaGraphDef` saved from this
session.
The exported `MetaGraphDef` will provide one `SignatureDef` for each
element of the export_outputs dict returned from the model_fn, named using
the same keys. One of these keys is always
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, indicating which
signature will be served when a serving request does not specify one.
For each signature, the outputs are provided by the corresponding
`ExportOutput`s, and the inputs are always the input receivers provided by
the serving_input_receiver_fn.
Extra assets may be written into the SavedModel via the extra_assets
argument. This should be a dict, where each key gives a destination path
(including the filename) relative to the assets.extra directory. The
corresponding value gives the full path of the source file to be copied.
For example, the simple case of copying a single file without renaming it
is specified as `{'my_asset_file.txt': '/path/to/my_asset_file.txt'}`.
Args:
export_dir_base: A string containing a directory in which to create
timestamped subdirectories containing exported SavedModels.
serving_input_receiver_fn: A function that takes no argument and
returns a `ServingInputReceiver`.
assets_extra: A dict specifying how to populate the assets.extra directory
within the exported SavedModel, or `None` if no extra assets are needed.
as_text: whether to write the SavedModel proto in text format.
checkpoint_path: The checkpoint path to export. If `None` (the default),
the most recent checkpoint found within the model directory is chosen.
Returns:
The string path to the exported directory.
Raises:
ValueError: if no serving_input_receiver_fn is provided, no export_outputs
are provided, or no checkpoint can be found.
"""
if serving_input_receiver_fn is None:
raise ValueError('serving_input_receiver_fn must be defined.')
with ops.Graph().as_default() as g:
self._create_and_assert_global_step(g)
random_seed.set_random_seed(self._config.tf_random_seed)
serving_input_receiver = serving_input_receiver_fn()
# Call the model_fn and collect the export_outputs.
estimator_spec = self._call_model_fn(
features=serving_input_receiver.features,
labels=None,
mode=model_fn_lib.ModeKeys.PREDICT,
config=self.config)
# Build the SignatureDefs from receivers and all outputs
signature_def_map = build_all_signature_defs(
serving_input_receiver.receiver_tensors,
estimator_spec.export_outputs,
serving_input_receiver.receiver_tensors_alternatives)
if not checkpoint_path:
# Locate the latest checkpoint
checkpoint_path = saver.latest_checkpoint(self._model_dir)
if not checkpoint_path:
raise ValueError("Couldn't find trained model at %s." % self._model_dir)
export_dir = get_timestamped_export_dir(export_dir_base)
temp_export_dir = get_temp_export_dir(export_dir)
# TODO(soergel): Consider whether MonitoredSession makes sense here
with tf_session.Session() as session:
saver_for_restore = estimator_spec.scaffold.saver or saver.Saver(
sharded=True)
saver_for_restore.restore(session, checkpoint_path)
# TODO(b/36111876): replace legacy_init_op with main_op mechanism
# pylint: disable=protected-access
local_init_op = (
estimator_spec.scaffold.local_init_op or
monitored_session.Scaffold._default_local_init_op())
# pylint: enable=protected-access
# Perform the export
builder = saved_model_builder.SavedModelBuilder(temp_export_dir)
builder.add_meta_graph_and_variables(
session, [tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=ops.get_collection(
ops.GraphKeys.ASSET_FILEPATHS),
legacy_init_op=local_init_op)
builder.save(as_text)
# Add the extra assets
if assets_extra:
assets_extra_path = os.path.join(compat.as_bytes(temp_export_dir),
compat.as_bytes('assets.extra'))
for dest_relative, source in assets_extra.items():
dest_absolute = os.path.join(compat.as_bytes(assets_extra_path),
compat.as_bytes(dest_relative))
dest_path = os.path.dirname(dest_absolute)
gfile.MakeDirs(dest_path)
gfile.Copy(source, dest_absolute)
gfile.Rename(temp_export_dir, export_dir)
return export_dir
def _add_meta_graph_for_mode(estimator,
builder,
input_receiver_fn,
checkpoint_path,
save_variables=True,
mode=model_fn_lib.ModeKeys.PREDICT,
export_tags=None,
check_variables=True):
if export_tags is None:
export_tags = model_fn_lib.EXPORT_TAG_MAP[mode]
with ops.Graph().as_default() as g:
random_seed.set_random_seed(estimator._config.tf_random_seed)
input_receiver = input_receiver_fn()
# Call the model_fn and collect the export_outputs.
estimator_spec = estimator._call_model_fn(
features=input_receiver.features,
labels=getattr(input_receiver, 'labels', None),
mode=mode,
config=estimator.config)
export_outputs = model_fn_lib.export_outputs_for_mode(
mode=estimator_spec.mode,
serving_export_outputs=estimator_spec.export_outputs,
predictions=estimator_spec.predictions,
loss=estimator_spec.loss,
metrics=estimator_spec.eval_metric_ops)
# Build the SignatureDefs from receivers and all outputs
signature_def_map = export_helpers.build_all_signature_defs(
input_receiver.receiver_tensors,
export_outputs,
getattr(input_receiver, 'receiver_tensors_alternatives', None),
serving_only=(mode == model_fn_lib.ModeKeys.PREDICT))
with tf_session.Session(config=estimator._session_config) as session:
if estimator_spec.scaffold.local_init_op is not None:
local_init_op = estimator_spec.scaffold.local_init_op
else:
local_init_op = monitored_session.Scaffold.default_local_init_op(
)
# This saver will be used both for restoring variables now,
# and in saving out the metagraph below. This ensures that any
# Custom Savers stored with the Scaffold are passed through to the
# SavedModel for restore later.
graph_saver = estimator_spec.scaffold.saver or saver.Saver(
sharded=True)
if save_variables and not check_variables:
raise ValueError(
'If `save_variables` is `True, `check_variables`'
'must not be `False`.')
if check_variables:
try:
graph_saver.restore(session, checkpoint_path)
except errors.NotFoundError as e:
msg = (
'Could not load all requested variables from checkpoint. '
'Please make sure your model_fn does not expect variables '
'that were not saved in the checkpoint.\n\n'
'Encountered error with mode `{}` while restoring '
'checkpoint from: `{}`. Full Traceback:\n\n{}').format(
mode, checkpoint_path, e)
raise ValueError(msg)
# We add the train op explicitly for now, so that we don't have to
# change the Builder public interface. Note that this is a no-op
# for prediction, where train_op is None.
builder._add_train_op(estimator_spec.train_op) # pylint: disable=protected-access
meta_graph_kwargs = dict(tags=export_tags,
signature_def_map=signature_def_map,
assets_collection=ops.get_collection(
ops.GraphKeys.ASSET_FILEPATHS),
strip_default_attrs=True,
legacy_init_op=local_init_op,
saver=graph_saver)
if save_variables:
builder.add_meta_graph_and_variables(session,
**meta_graph_kwargs)
else:
builder.add_meta_graph(**meta_graph_kwargs)
