Merge pull request tensorflow#4606 from jhseu/branch_134473452

Branch 134473452

Author: jhseu

tensorflow/BUILD

@@ -113,7 +113,6 @@ filegroup(
         "//tensorflow/contrib/learn:all_files",
         "//tensorflow/contrib/learn/python/learn/datasets:all_files",
         "//tensorflow/contrib/linear_optimizer:all_files",
-        "//tensorflow/contrib/linear_optimizer/kernels:all_files",
         "//tensorflow/contrib/lookup:all_files",
         "//tensorflow/contrib/losses:all_files",
         "//tensorflow/contrib/metrics:all_files",

tensorflow/c/c_api.cc

@@ -1228,10 +1228,10 @@ int TF_OperationGetControlOutputs(TF_Operation* oper,
   return count;
 }
 
-TF_Attr_Metadata TF_OperationGetAttrMetadata(TF_Operation* oper,
-                                             const char* attr_name,
-                                             TF_Status* status) {
-  TF_Attr_Metadata metadata;
+TF_AttrMetadata TF_OperationGetAttrMetadata(TF_Operation* oper,
+                                            const char* attr_name,
+                                            TF_Status* status) {
+  TF_AttrMetadata metadata;
   const auto* attr = GetAttrValue(oper, attr_name, status);
   if (!status->status.ok()) return metadata;
   switch (attr->value_case()) {

tensorflow/c/c_api.h

@@ -537,7 +537,7 @@ extern int TF_OperationGetControlOutputs(TF_Operation* oper,
                                          TF_Operation** control_outputs,
                                          int max_control_outputs);
 
-// TF_Attr_Type describes the type of the value of an attribute on an operation.
+// TF_AttrType describes the type of the value of an attribute on an operation.
 typedef enum {
   TF_ATTR_STRING = 0,
   TF_ATTR_INT = 1,
@@ -548,9 +548,9 @@ typedef enum {
   TF_ATTR_TENSOR = 6,
   TF_ATTR_PLACEHOLDER = 7,
   TF_ATTR_FUNC = 8,
-} TF_Attr_Type;
+} TF_AttrType;
 
-// TF_Attr_Metadata describes the value of an attribute on an operation.
+// TF_AttrMetadata describes the value of an attribute on an operation.
 typedef struct {
   // A boolean: 1 if the attribute value is a list, 0 otherwise.
   unsigned char is_list;
@@ -560,7 +560,7 @@ typedef struct {
 
   // Type of elements of the list if is_list != 0.
   // Type of the single value stored in the attribute if is_list == 0.
-  TF_Attr_Type type;
+  TF_AttrType type;
 
   // Total size the attribute value.
   // The units of total_size depend on is_list and type.
@@ -579,16 +579,16 @@ typedef struct {
   //     of dimensions of all shapes in the list.
   // (5) Otherwise, total_size is undefined.
   int64_t total_size;
-} TF_Attr_Metadata;
+} TF_AttrMetadata;
 
 // Returns metadata about the value of the attribute `attr_name` of `oper`.
-TF_Attr_Metadata TF_OperationGetAttrMetadata(TF_Operation* oper,
-                                             const char* attr_name,
-                                             TF_Status* status);
+TF_AttrMetadata TF_OperationGetAttrMetadata(TF_Operation* oper,
+                                            const char* attr_name,
+                                            TF_Status* status);
 
 // Fills in `value` with the value of the attribute `attr_name`.  `value` must
 // point to an array of length at least `max_length` (ideally set to
-// TF_Attr_Metadata.total_size from TF_OperationGetAttrMetadata(oper,
+// TF_AttrMetadata.total_size from TF_OperationGetAttrMetadata(oper,
 // attr_name)).
 extern void TF_OperationGetAttrString(TF_Operation* oper, const char* attr_name,
                                       void* value, int max_length,
@@ -600,8 +600,8 @@ extern void TF_OperationGetAttrString(TF_Operation* oper, const char* attr_name,
 //
 // The elements of values will point to addresses in `storage` which must be at
 // least `storage_size` bytes large.  Ideally, max_values would be set to
-// TF_Attr_Metadata.list_size and `storage` would be at least
-// TF_Attr_Metadata.total_size, obtained from TF_OperationGetAttrMetadata(oper,
+// TF_AttrMetadata.list_size and `storage` would be at least
+// TF_AttrMetadata.total_size, obtained from TF_OperationGetAttrMetadata(oper,
 // attr_name).
 //
 // Fails if storage_size is too small to hold the requested number of strings.
@@ -616,7 +616,7 @@ extern void TF_OperationGetAttrInt(TF_Operation* oper, const char* attr_name,
 
 // Fills in `values` with the value of the attribute `attr_name` of `oper`.
 // `values` must point to an array of length at least `max_values` (ideally set
-// TF_Attr_Metadata.list_size from TF_OperationGetAttrMetadata(oper,
+// TF_AttrMetadata.list_size from TF_OperationGetAttrMetadata(oper,
 // attr_name)).
 extern void TF_OperationGetAttrIntList(TF_Operation* oper,
                                        const char* attr_name, int64_t* values,
@@ -627,7 +627,7 @@ extern void TF_OperationGetAttrFloat(TF_Operation* oper, const char* attr_name,
 
 // Fills in `values` with the value of the attribute `attr_name` of `oper`.
 // `values` must point to an array of length at least `max_values` (ideally set
-// to TF_Attr_Metadata.list_size from TF_OperationGetAttrMetadata(oper,
+// to TF_AttrMetadata.list_size from TF_OperationGetAttrMetadata(oper,
 // attr_name)).
 extern void TF_OperationGetAttrFloatList(TF_Operation* oper,
                                          const char* attr_name, float* values,
@@ -638,7 +638,7 @@ extern void TF_OperationGetAttrBool(TF_Operation* oper, const char* attr_name,
 
 // Fills in `values` with the value of the attribute `attr_name` of `oper`.
 // `values` must point to an array of length at least `max_values` (ideally set
-// to TF_Attr_Metadata.list_size from TF_OperationGetAttrMetadata(oper,
+// to TF_AttrMetadata.list_size from TF_OperationGetAttrMetadata(oper,
 // attr_name)).
 extern void TF_OperationGetAttrBoolList(TF_Operation* oper,
                                         const char* attr_name,
@@ -650,7 +650,7 @@ extern void TF_OperationGetAttrType(TF_Operation* oper, const char* attr_name,
 
 // Fills in `values` with the value of the attribute `attr_name` of `oper`.
 // `values` must point to an array of length at least `max_values` (ideally set
-// to TF_Attr_Metadata.list_size from TF_OperationGetAttrMetadata(oper,
+// to TF_AttrMetadata.list_size from TF_OperationGetAttrMetadata(oper,
 // attr_name)).
 extern void TF_OperationGetAttrTypeList(TF_Operation* oper,
                                         const char* attr_name,
@@ -672,8 +672,8 @@ extern void TF_OperationGetAttrShape(TF_Operation* oper, const char* attr_name,
 //
 // The elements of `dims` will point to addresses in `storage` which must be
 // large enough to hold at least `storage_size` int64_ts.  Ideally, `num_shapes`
-// would be set to TF_Attr_Metadata.list_size and `storage_size` would be set to
-// TF_Attr_Metadata.total_size from TF_OperationGetAttrMetadata(oper,
+// would be set to TF_AttrMetadata.list_size and `storage_size` would be set to
+// TF_AttrMetadata.total_size from TF_OperationGetAttrMetadata(oper,
 // attr_name).
 //
 // Fails if storage_size is insufficient to hold the requested shapes.
@@ -692,7 +692,7 @@ extern void TF_OperationGetAttrTensorShapeProto(TF_Operation* oper,
 
 // Fills in `values` with binary-serialized TensorShapeProto values of the
 // attribute `attr_name` of `oper`. `values` must point to an array of length at
-// least `num_values` (ideally set to TF_Attr_Metadata.list_size from
+// least `num_values` (ideally set to TF_AttrMetadata.list_size from
 // TF_OperationGetAttrMetadata(oper, attr_name)).
 extern void TF_OperationGetAttrTensorShapeProtoList(TF_Operation* oper,
                                                     const char* attr_name,
@@ -709,7 +709,7 @@ extern void TF_OperationGetAttrTensor(TF_Operation* oper, const char* attr_name,
 
 // Fills in `values` with the TF_Tensor values of the attribute `attr_name` of
 // `oper`. `values` must point to an array of TF_Tensor* of length at least
-// `max_values` (ideally set to TF_Attr_Metadata.list_size from
+// `max_values` (ideally set to TF_AttrMetadata.list_size from
 // TF_OperationGetAttrMetadata(oper, attr_name)).
 //
 // The caller takes ownership of all the non-null TF_Tensor* entries in `values`

tensorflow/c/c_api_test.cc

@@ -864,7 +864,7 @@ TEST(CAPI, ColocateWith) {
   TF_Operation* add = TF_FinishOperation(desc, s);
   ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
 
-  TF_Attr_Metadata m =
+  TF_AttrMetadata m =
       TF_OperationGetAttrMetadata(add, tensorflow::kColocationAttrName, s);
   EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
   EXPECT_EQ(1, m.is_list);

tensorflow/cc/saved_model/BUILD

@@ -52,6 +52,7 @@ filegroup(
     srcs = glob([
         "testdata/half_plus_two/**",
         "testdata/half_plus_two_pbtxt/**",
+        "testdata/half_plus_two_sharded/**",
     ]),
 )
 

tensorflow/cc/saved_model/constants.h

@@ -30,6 +30,10 @@ constexpr char kSavedModelVariablesDirectory[] = "variables";
 // SavedModel variables filename.
 constexpr char kSavedModelVariablesFilename[] = "saved_model_variables";
 
+// SavedModel sharded variables filename.
+constexpr char kSavedModelVariablesShardedFilename[] =
+    "saved_model_variables-\?\?\?\?\?-of-\?\?\?\?\?";
+
 // Commonly used tags.
 constexpr char kSavedModelTagServe[] = "serve";
 constexpr char kSavedModelTagTrain[] = "train";

tensorflow/cc/saved_model/loader.cc

@@ -86,14 +86,18 @@ Status Restore(const RunOptions& run_options, const string& export_dir,
                const StringPiece restore_op_name,
                const StringPiece variable_filename_const_op_name,
                Session* session) {
-  const string variables_path = io::JoinPath(
-      export_dir, kSavedModelVariablesDirectory, kSavedModelVariablesFilename);
-  if (!Env::Default()->FileExists(variables_path)) {
-    return Status(
-        error::Code::NOT_FOUND,
-        "Could not find checkpointed variables at: " + variables_path);
+  // Find path to variables to be restored in export directory.
+  string variables_path =
+      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;
   }
-
   // Add variables to the graph.
   Tensor variables_path_tensor(DT_STRING, TensorShape({}));
   variables_path_tensor.scalar<string>()() = variables_path;

tensorflow/cc/saved_model/loader_test.cc

@@ -28,6 +28,8 @@ namespace {
 
 constexpr char kTestDataPb[] = "cc/saved_model/testdata/half_plus_two";
 constexpr char kTestDataPbTxt[] = "cc/saved_model/testdata/half_plus_two_pbtxt";
+constexpr char kTestDataSharded[] =
+    "cc/saved_model/testdata/half_plus_two_sharded";
 
 class LoaderTest : public ::testing::Test {
  protected:
@@ -110,6 +112,18 @@ TEST_F(LoaderTest, PbtxtFormat) {
   CheckSavedModelBundle(bundle);
 }
 
+TEST_F(LoaderTest, ShardedVariables) {
+  SavedModelBundle bundle;
+  SessionOptions session_options;
+  RunOptions run_options;
+
+  const string export_dir =
+      io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
+  TF_ASSERT_OK(LoadSavedModel(export_dir, {kSavedModelTagServe},
+                              session_options, run_options, &bundle));
+  CheckSavedModelBundle(bundle);
+}
+
 TEST_F(LoaderTest, InvalidExportPath) {
   SavedModelBundle bundle;
   RunOptions run_options;

tensorflow/cc/saved_model/testdata/half_plus_two_sharded/saved_model.pb

@@ -0,0 +1,2 @@
+model_checkpoint_path: "/tmp/saved_model/half_plus_two/variables/saved_model_variables-?????-of-00001"
+all_model_checkpoint_paths: "/tmp/saved_model/half_plus_two/variables/saved_model_variables-?????-of-00001"