From 5bb0f5174d24ef76c5b58929f70f1f5fc7d82e67 Mon Sep 17 00:00:00 2001
From: Cesar Crusius <ccrusius@google.com>
Date: Mon, 25 Jan 2021 10:49:29 -0800
Subject: [PATCH] Add an option to save only the "TF Program" section of a
 Saved Model.

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
---
 tensorflow/python/saved_model/save.py | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/tensorflow/python/saved_model/save.py b/tensorflow/python/saved_model/save.py
index 99ed53140ac..1979c92c980 100644
--- a/tensorflow/python/saved_model/save.py
+++ b/tensorflow/python/saved_model/save.py
@@ -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.