experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add an option to save only the "TF Program" section of a Saved Model.

Browse Source 
This allows for saving TF programs without saving the checkpoint (the
"state"). This option will generate saved models that are not fully
supported under `saved_model.load`, but that can be used in
experimental settings where the loading steps are under full user
control.

PiperOrigin-RevId: 353681612
Change-Id: Idada850426d0ab728cdf1bda3b8a958f9239223a
This commit is contained in:
Cesar Crusius 2021-01-25 10:49:29 -08:00 committed by TensorFlower Gardener
parent fbf6fed2c7
commit 5bb0f5174d
1 changed files with 16 additions and 9 deletions

25
tensorflow/python/saved_model/save.py
View File

@ -1048,8 +1048,12 @@ def save(obj, export_dir, signatures=None, options=None):
raise_metadata_warning=True)
def save_and_return_nodes(obj, export_dir, signatures=None, options=None,
raise_metadata_warning=False):
def save_and_return_nodes(obj,
export_dir,
signatures=None,
options=None,
raise_metadata_warning=False,
experimental_skip_checkpoint=False):
"""Saves a SavedModel while returning all saved nodes and their paths.
Please see `tf.saved_model.save` for details.
@ -1062,6 +1066,8 @@ def save_and_return_nodes(obj, export_dir, signatures=None, options=None,
options: `tf.saved_model.SaveOptions` object for configuring save options.
raise_metadata_warning: Whether to raise the metadata warning. This arg will
be removed in TF 2.5.
experimental_skip_checkpoint: If set to `True`, the checkpoint will not
be written.
Returns:
A tuple of (a list of saved nodes in the order they are serialized to the
@ -1082,13 +1088,14 @@ def save_and_return_nodes(obj, export_dir, signatures=None, options=None,
# Write the checkpoint, copy assets into the assets directory, and write out
# the SavedModel proto itself.
utils_impl.get_or_create_variables_dir(export_dir)
ckpt_options = checkpoint_options.CheckpointOptions(
experimental_io_device=options.experimental_io_device)
object_saver.save(
utils_impl.get_variables_path(export_dir), options=ckpt_options)
builder_impl.copy_assets_to_destination_dir(asset_info.asset_filename_map,
export_dir)
if not experimental_skip_checkpoint:
utils_impl.get_or_create_variables_dir(export_dir)
ckpt_options = checkpoint_options.CheckpointOptions(
experimental_io_device=options.experimental_io_device)
object_saver.save(
utils_impl.get_variables_path(export_dir), options=ckpt_options)
builder_impl.copy_assets_to_destination_dir(asset_info.asset_filename_map,
export_dir)
# Note that this needs to be the last file operation when saving the
# SavedModel. Users rely on checking saved_model_dir/saved_model.pb as an
# indication that the SavedModel is completely written.