diff --git a/tensorflow/cc/saved_model/BUILD b/tensorflow/cc/saved_model/BUILD index bb2a6063b5c..6a55ecf9369 100644 --- a/tensorflow/cc/saved_model/BUILD +++ b/tensorflow/cc/saved_model/BUILD @@ -31,6 +31,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:tensorflow", + "//tensorflow/core/util/tensor_bundle:naming", ], ) diff --git a/tensorflow/cc/saved_model/constants.h b/tensorflow/cc/saved_model/constants.h index b97f6c84faf..61ffa7a09ac 100644 --- a/tensorflow/cc/saved_model/constants.h +++ b/tensorflow/cc/saved_model/constants.h @@ -36,10 +36,6 @@ constexpr char kSavedModelVariablesDirectory[] = "variables"; // SavedModel variables filename. constexpr char kSavedModelVariablesFilename[] = "variables"; -// SavedModel sharded variables filename. -constexpr char kSavedModelVariablesShardedFilename[] = - "variables-\?\?\?\?\?-of-\?\?\?\?\?"; - // Commonly used tags. constexpr char kSavedModelTagServe[] = "serve"; constexpr char kSavedModelTagTrain[] = "train"; diff --git a/tensorflow/cc/saved_model/loader.cc b/tensorflow/cc/saved_model/loader.cc index 329bea48fa8..1f952293550 100644 --- a/tensorflow/cc/saved_model/loader.cc +++ b/tensorflow/cc/saved_model/loader.cc @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/core/protobuf/saved_model.pb.h" #include "tensorflow/core/public/session.h" #include "tensorflow/core/public/session_options.h" +#include "tensorflow/core/util/tensor_bundle/naming.h" namespace tensorflow { namespace { @@ -87,17 +88,20 @@ Status Restore(const RunOptions& run_options, const string& export_dir, const StringPiece variable_filename_const_op_name, Session* session) { // Find path to variables to be restored in export directory. - string variables_path = + const string variables_directory = io::JoinPath(export_dir, kSavedModelVariablesDirectory); - const string unsharded_variables_path = - io::JoinPath(variables_path, kSavedModelVariablesFilename); - if (Env::Default()->FileExists(unsharded_variables_path)) { - variables_path = unsharded_variables_path; - } else { - const string sharded_variables_path = - io::JoinPath(variables_path, kSavedModelVariablesShardedFilename); - variables_path = sharded_variables_path; + // Check for saver checkpoints in v2 format. Models exported in the checkpoint + // v2 format will have a variables.index file. The corresponding + // variables are stored in the variables.data-?????-of-????? files. + const string variables_index_path = io::JoinPath( + variables_directory, MetaFilename(kSavedModelVariablesFilename)); + if (!Env::Default()->FileExists(variables_index_path)) { + return errors::NotFound( + "Checkpoint index file not found in SavedModel directory."); } + const string variables_path = + io::JoinPath(variables_directory, kSavedModelVariablesFilename); + // Add variables to the graph. Tensor variables_path_tensor(DT_STRING, TensorShape({})); variables_path_tensor.scalar()() = variables_path; diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two/saved_model.pb b/tensorflow/cc/saved_model/testdata/half_plus_two/saved_model.pb index d0f0853aa87..e894f9b1011 100644 Binary files a/tensorflow/cc/saved_model/testdata/half_plus_two/saved_model.pb and b/tensorflow/cc/saved_model/testdata/half_plus_two/saved_model.pb differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two/variables/variables b/tensorflow/cc/saved_model/testdata/half_plus_two/variables/variables deleted file mode 100644 index e1ac9e900e8..00000000000 Binary files a/tensorflow/cc/saved_model/testdata/half_plus_two/variables/variables and /dev/null differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two/variables/variables.data-00000-of-00001 b/tensorflow/cc/saved_model/testdata/half_plus_two/variables/variables.data-00000-of-00001 new file mode 100644 index 00000000000..20bc7d454dd Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/half_plus_two/variables/variables.data-00000-of-00001 differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two/variables/variables.index b/tensorflow/cc/saved_model/testdata/half_plus_two/variables/variables.index new file mode 100644 index 00000000000..e7df518f5b5 Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/half_plus_two/variables/variables.index differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/variables/variables-00000-of-00001 b/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/variables/variables-00000-of-00001 deleted file mode 100644 index e1ac9e900e8..00000000000 Binary files a/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/variables/variables-00000-of-00001 and /dev/null differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/variables/variables.data-00000-of-00001 b/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/variables/variables.data-00000-of-00001 new file mode 100644 index 00000000000..20bc7d454dd Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/variables/variables.data-00000-of-00001 differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/variables/variables.index b/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/variables/variables.index new file mode 100644 index 00000000000..e7df518f5b5 Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/variables/variables.index differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/saved_model.pb b/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/saved_model.pb index d0f0853aa87..e894f9b1011 100644 Binary files a/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/saved_model.pb and b/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/saved_model.pb differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/variables/variables-00000-of-00001 b/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/variables/variables-00000-of-00001 deleted file mode 100644 index e1ac9e900e8..00000000000 Binary files a/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/variables/variables-00000-of-00001 and /dev/null differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/variables/variables.data-00000-of-00001 b/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/variables/variables.data-00000-of-00001 new file mode 100644 index 00000000000..20bc7d454dd Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/variables/variables.data-00000-of-00001 differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/variables/variables.index b/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/variables/variables.index new file mode 100644 index 00000000000..e7df518f5b5 Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/half_plus_two_sharded/variables/variables.index differ diff --git a/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/saved_model.pb b/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/saved_model.pb index d0f0853aa87..e894f9b1011 100644 Binary files a/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/saved_model.pb and b/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/saved_model.pb differ diff --git a/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/checkpoint b/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/checkpoint deleted file mode 100644 index 1065013315f..00000000000 --- a/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/checkpoint +++ /dev/null @@ -1,2 +0,0 @@ -model_checkpoint_path: "/tmp/saved_model/half_plus_two/variables/variables-?????-of-00001" -all_model_checkpoint_paths: "/tmp/saved_model/half_plus_two/variables/variables-?????-of-00001" diff --git a/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/variables-00000-of-00001 b/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/variables-00000-of-00001 deleted file mode 100644 index e1ac9e900e8..00000000000 Binary files a/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/variables-00000-of-00001 and /dev/null differ diff --git a/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/variables.data-00000-of-00001 b/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/variables.data-00000-of-00001 new file mode 100644 index 00000000000..20bc7d454dd Binary files /dev/null and b/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/variables.data-00000-of-00001 differ diff --git a/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/variables.index b/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/variables.index new file mode 100644 index 00000000000..e7df518f5b5 Binary files /dev/null and b/tensorflow/contrib/session_bundle/testdata/saved_model_half_plus_two/variables/variables.index differ diff --git a/tensorflow/python/saved_model/builder.py b/tensorflow/python/saved_model/builder.py index 4f3e8e138e8..1e28b979f73 100644 --- a/tensorflow/python/saved_model/builder.py +++ b/tensorflow/python/saved_model/builder.py @@ -256,7 +256,7 @@ class SavedModelBuilder(object): saver = tf_saver.Saver( variables.all_variables(), sharded=True, - write_version=saver_pb2.SaverDef.V1) + write_version=saver_pb2.SaverDef.V2) meta_graph_def = saver.export_meta_graph() # Tag the meta graph def and add it to the SavedModel. @@ -305,7 +305,7 @@ class SavedModelBuilder(object): saver = tf_saver.Saver( variables.all_variables(), sharded=True, - write_version=saver_pb2.SaverDef.V1) + write_version=saver_pb2.SaverDef.V2) saver.save(sess, variables_path, write_meta_graph=False) meta_graph_def = saver.export_meta_graph() diff --git a/tensorflow/python/saved_model/constants.py b/tensorflow/python/saved_model/constants.py index d88387bab92..b7b22dda4ef 100644 --- a/tensorflow/python/saved_model/constants.py +++ b/tensorflow/python/saved_model/constants.py @@ -31,4 +31,3 @@ TAG_TRAINING = "train" VARIABLES_DIRECTORY = "variables" VARIABLES_FILENAME = "variables" -VARIABLES_FILENAME_SHARDED = VARIABLES_FILENAME + "-?????-of-?????" diff --git a/tensorflow/python/saved_model/example/saved_model_half_plus_two.py b/tensorflow/python/saved_model/example/saved_model_half_plus_two.py index 71cff17c6ac..76175e86f01 100644 --- a/tensorflow/python/saved_model/example/saved_model_half_plus_two.py +++ b/tensorflow/python/saved_model/example/saved_model_half_plus_two.py @@ -124,11 +124,11 @@ def _generate_saved_model_for_half_plus_two(export_dir, as_text=False): def main(_): - export_dir_pb = "/tmp/saved_model/half_plus_two" + export_dir_pb = "/tmp/saved_model/v2_half_plus_two_unsharded" _generate_saved_model_for_half_plus_two(export_dir_pb) print("SavedModel generated at: %s" % export_dir_pb) - export_dir_pbtxt = "/tmp/saved_model/half_plus_two_pbtxt" + export_dir_pbtxt = "/tmp/saved_model/v2_half_plus_two_pbtxt" _generate_saved_model_for_half_plus_two(export_dir_pbtxt, as_text=True) print("SavedModel generated at: %s" % export_dir_pbtxt) diff --git a/tensorflow/python/saved_model/loader.py b/tensorflow/python/saved_model/loader.py index 594416287da..451e0748c75 100644 --- a/tensorflow/python/saved_model/loader.py +++ b/tensorflow/python/saved_model/loader.py @@ -188,7 +188,7 @@ def load(sess, tags, export_dir): variables_path = os.path.join( compat.as_bytes(export_dir), compat.as_bytes(constants.VARIABLES_DIRECTORY), - compat.as_bytes(constants.VARIABLES_FILENAME_SHARDED)) + compat.as_bytes(constants.VARIABLES_FILENAME)) # Restore the variables using the built saver in the provided session. saver.restore(sess, variables_path)