Update SavedModel loader tests to use constants and serialized tf Example.

Change: 135944464

5 files changed, 529 insertions, 67 deletions

diff --git a/tensorflow/cc/saved_model/BUILD b/tensorflow/cc/saved_model/BUILD
index e0cdf06938..bb2a6063b5 100644
--- a/tensorflow/cc/saved_model/BUILD
+++ b/tensorflow/cc/saved_model/BUILD
@@ -28,7 +28,6 @@ cc_library(
     deps = [
         ":constants",
         "//tensorflow/core:core_cpu",
-        "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:tensorflow",
@@ -45,7 +44,9 @@ tf_cc_test(
     deps = [
         ":constants",
         ":loader",
+        ":signature_constants",
         "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
diff --git a/tensorflow/cc/saved_model/loader_test.cc b/tensorflow/cc/saved_model/loader_test.cc
index 817fd8fa95..6011531f0d 100644
--- a/tensorflow/cc/saved_model/loader_test.cc
+++ b/tensorflow/cc/saved_model/loader_test.cc
@@ -16,6 +16,9 @@ limitations under the License.
 #include "tensorflow/cc/saved_model/loader.h"
 
 #include "tensorflow/cc/saved_model/constants.h"
+#include "tensorflow/cc/saved_model/signature_constants.h"
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/example/feature.pb.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
@@ -34,17 +37,30 @@ class LoaderTest : public ::testing::Test {
  protected:
   LoaderTest() {}
 
-  void CheckSavedModelBundle(const SavedModelBundle& bundle) {
-    // Validate the half plus two behavior.
-    Tensor input = test::AsTensor<float>({0, 1, 2, 3}, TensorShape({4, 1}));
+  string MakeSerializedExample(float x) {
+    tensorflow::Example example;
+    auto* feature_map = example.mutable_features()->mutable_feature();
+    (*feature_map)["x"].mutable_float_list()->add_value(x);
+    return example.SerializeAsString();
+  }
 
+  void CheckSavedModelBundle(const SavedModelBundle& bundle) {
     // Retrieve the regression signature from meta graph def.
     const auto signature_def_map = bundle.meta_graph_def.signature_def();
-    const auto signature_def = signature_def_map.at("regression");
+    const auto signature_def = signature_def_map.at(kRegressMethodName);
+
+    const string input_name = signature_def.inputs().at(kRegressInputs).name();
+    const string output_name =
+        signature_def.outputs().at(kRegressOutputs).name();
 
-    const string input_name = signature_def.inputs().at("input").name();
-    const string output_name = signature_def.outputs().at("output").name();
+    std::vector<string> serialized_examples;
+    for (float x : {0, 1, 2, 3}) {
+      serialized_examples.push_back(MakeSerializedExample(x));
+    }
 
+    // Validate the half plus two behavior.
+    Tensor input =
+        test::AsTensor<string>(serialized_examples, TensorShape({4}));
     std::vector<Tensor> outputs;
     TF_ASSERT_OK(bundle.session->Run({{input_name, input}}, {output_name}, {},
                                      &outputs));
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 5a2dd4dd84..d0f0853aa8 100644
--- a/tensorflow/cc/saved_model/testdata/half_plus_two/saved_model.pb
+++ b/tensorflow/cc/saved_model/testdata/half_plus_two/saved_model.pb
diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/saved_model.pbtxt b/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/saved_model.pbtxt
index 30c2c25a19..2e714d262d 100644
--- a/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/saved_model.pbtxt
+++ b/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/saved_model.pbtxt
@@ -141,6 +141,88 @@ meta_graphs {
         name: "NoOp"
       }
       op {
+        name: "ParseExample"
+        input_arg {
+          name: "serialized"
+          type: DT_STRING
+        }
+        input_arg {
+          name: "names"
+          type: DT_STRING
+        }
+        input_arg {
+          name: "sparse_keys"
+          type: DT_STRING
+          number_attr: "Nsparse"
+        }
+        input_arg {
+          name: "dense_keys"
+          type: DT_STRING
+          number_attr: "Ndense"
+        }
+        input_arg {
+          name: "dense_defaults"
+          type_list_attr: "Tdense"
+        }
+        output_arg {
+          name: "sparse_indices"
+          type: DT_INT64
+          number_attr: "Nsparse"
+        }
+        output_arg {
+          name: "sparse_values"
+          type_list_attr: "sparse_types"
+        }
+        output_arg {
+          name: "sparse_shapes"
+          type: DT_INT64
+          number_attr: "Nsparse"
+        }
+        output_arg {
+          name: "dense_values"
+          type_list_attr: "Tdense"
+        }
+        attr {
+          name: "Nsparse"
+          type: "int"
+          has_minimum: true
+        }
+        attr {
+          name: "Ndense"
+          type: "int"
+          has_minimum: true
+        }
+        attr {
+          name: "sparse_types"
+          type: "list(type)"
+          has_minimum: true
+          allowed_values {
+            list {
+              type: DT_FLOAT
+              type: DT_INT64
+              type: DT_STRING
+            }
+          }
+        }
+        attr {
+          name: "Tdense"
+          type: "list(type)"
+          has_minimum: true
+          allowed_values {
+            list {
+              type: DT_FLOAT
+              type: DT_INT64
+              type: DT_STRING
+            }
+          }
+        }
+        attr {
+          name: "dense_shapes"
+          type: "list(shape)"
+          has_minimum: true
+        }
+      }
+      op {
         name: "Placeholder"
         output_arg {
           name: "output"
@@ -160,33 +242,28 @@ meta_graphs {
         }
       }
       op {
-        name: "RestoreSlice"
+        name: "RestoreV2"
         input_arg {
-          name: "file_pattern"
+          name: "prefix"
           type: DT_STRING
         }
         input_arg {
-          name: "tensor_name"
+          name: "tensor_names"
           type: DT_STRING
         }
         input_arg {
-          name: "shape_and_slice"
+          name: "shape_and_slices"
           type: DT_STRING
         }
         output_arg {
-          name: "tensor"
-          type_attr: "dt"
+          name: "tensors"
+          type_list_attr: "dtypes"
         }
         attr {
-          name: "dt"
-          type: "type"
-        }
-        attr {
-          name: "preferred_shard"
-          type: "int"
-          default_value {
-            i: -1
-          }
+          name: "dtypes"
+          type: "list(type)"
+          has_minimum: true
+          minimum: 1
         }
       }
       op {
@@ -215,6 +292,40 @@ meta_graphs {
         }
       }
       op {
+        name: "ShardedFilename"
+        input_arg {
+          name: "basename"
+          type: DT_STRING
+        }
+        input_arg {
+          name: "shard"
+          type: DT_INT32
+        }
+        input_arg {
+          name: "num_shards"
+          type: DT_INT32
+        }
+        output_arg {
+          name: "filename"
+          type: DT_STRING
+        }
+      }
+      op {
+        name: "ShardedFilespec"
+        input_arg {
+          name: "basename"
+          type: DT_STRING
+        }
+        input_arg {
+          name: "num_shards"
+          type: DT_INT32
+        }
+        output_arg {
+          name: "filename"
+          type: DT_STRING
+        }
+      }
+      op {
         name: "Variable"
         output_arg {
           name: "ref"
@@ -524,7 +635,7 @@ meta_graphs {
       }
     }
     node {
-      name: "x"
+      name: "tf_example"
       op: "Placeholder"
       attr {
         key: "_output_shapes"
@@ -539,7 +650,7 @@ meta_graphs {
       attr {
         key: "dtype"
         value {
-          type: DT_FLOAT
+          type: DT_STRING
         }
       }
       attr {
@@ -551,6 +662,190 @@ meta_graphs {
       }
     }
     node {
+      name: "ParseExample/Const"
+      op: "Const"
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+              dim {
+              }
+            }
+          }
+        }
+      }
+      attr {
+        key: "dtype"
+        value {
+          type: DT_FLOAT
+        }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+              }
+            }
+          }
+        }
+      }
+    }
+    node {
+      name: "ParseExample/ParseExample/names"
+      op: "Const"
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+              dim {
+              }
+            }
+          }
+        }
+      }
+      attr {
+        key: "dtype"
+        value {
+          type: DT_STRING
+        }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_STRING
+            tensor_shape {
+              dim {
+              }
+            }
+          }
+        }
+      }
+    }
+    node {
+      name: "ParseExample/ParseExample/dense_keys_0"
+      op: "Const"
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+            }
+          }
+        }
+      }
+      attr {
+        key: "dtype"
+        value {
+          type: DT_STRING
+        }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_STRING
+            tensor_shape {
+            }
+            string_val: "x"
+          }
+        }
+      }
+    }
+    node {
+      name: "ParseExample/ParseExample"
+      op: "ParseExample"
+      input: "tf_example"
+      input: "ParseExample/ParseExample/names"
+      input: "ParseExample/ParseExample/dense_keys_0"
+      input: "ParseExample/Const"
+      attr {
+        key: "Ndense"
+        value {
+          i: 1
+        }
+      }
+      attr {
+        key: "Nsparse"
+        value {
+          i: 0
+        }
+      }
+      attr {
+        key: "Tdense"
+        value {
+          list {
+            type: DT_FLOAT
+          }
+        }
+      }
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+              dim {
+                size: -1
+              }
+              dim {
+                size: 1
+              }
+            }
+          }
+        }
+      }
+      attr {
+        key: "dense_shapes"
+        value {
+          list {
+            shape {
+              dim {
+                size: 1
+              }
+            }
+          }
+        }
+      }
+      attr {
+        key: "sparse_types"
+        value {
+          list {
+          }
+        }
+      }
+    }
+    node {
+      name: "x"
+      op: "Identity"
+      input: "ParseExample/ParseExample"
+      attr {
+        key: "T"
+        value {
+          type: DT_FLOAT
+        }
+      }
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+              dim {
+                size: -1
+              }
+              dim {
+                size: 1
+              }
+            }
+          }
+        }
+      }
+    }
+    node {
       name: "Mul"
       op: "Mul"
       input: "a/read"
@@ -566,7 +861,12 @@ meta_graphs {
         value {
           list {
             shape {
-              unknown_rank: true
+              dim {
+                size: -1
+              }
+              dim {
+                size: 1
+              }
             }
           }
         }
@@ -588,7 +888,38 @@ meta_graphs {
         value {
           list {
             shape {
-              unknown_rank: true
+              dim {
+                size: -1
+              }
+              dim {
+                size: 1
+              }
+            }
+          }
+        }
+      }
+    }
+    node {
+      name: "Identity"
+      op: "Identity"
+      input: "y"
+      attr {
+        key: "T"
+        value {
+          type: DT_FLOAT
+        }
+      }
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+              dim {
+                size: -1
+              }
+              dim {
+                size: 1
+              }
             }
           }
         }
@@ -631,6 +962,82 @@ meta_graphs {
       }
     }
     node {
+      name: "save/num_shards"
+      op: "Const"
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+            }
+          }
+        }
+      }
+      attr {
+        key: "dtype"
+        value {
+          type: DT_INT32
+        }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_INT32
+            tensor_shape {
+            }
+            int_val: 1
+          }
+        }
+      }
+    }
+    node {
+      name: "save/ShardedFilename/shard"
+      op: "Const"
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+            }
+          }
+        }
+      }
+      attr {
+        key: "dtype"
+        value {
+          type: DT_INT32
+        }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_INT32
+            tensor_shape {
+            }
+            int_val: 0
+          }
+        }
+      }
+    }
+    node {
+      name: "save/ShardedFilename"
+      op: "ShardedFilename"
+      input: "save/Const"
+      input: "save/ShardedFilename/shard"
+      input: "save/num_shards"
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+            }
+          }
+        }
+      }
+    }
+    node {
       name: "save/save/tensor_names"
       op: "Const"
       attr {
@@ -707,7 +1114,7 @@ meta_graphs {
     node {
       name: "save/save"
       op: "SaveSlices"
-      input: "save/Const"
+      input: "save/ShardedFilename"
       input: "save/save/tensor_names"
       input: "save/save/shapes_and_slices"
       input: "a"
@@ -725,7 +1132,7 @@ meta_graphs {
     node {
       name: "save/control_dependency"
       op: "Identity"
-      input: "save/Const"
+      input: "save/ShardedFilename"
       input: "^save/save"
       attr {
         key: "T"
@@ -737,7 +1144,7 @@ meta_graphs {
         key: "_class"
         value {
           list {
-            s: "loc:@save/Const"
+            s: "loc:@save/ShardedFilename"
           }
         }
       }
@@ -752,13 +1159,32 @@ meta_graphs {
       }
     }
     node {
-      name: "save/restore_slice/tensor_name"
+      name: "save/ShardedFilespec"
+      op: "ShardedFilespec"
+      input: "save/Const"
+      input: "save/num_shards"
+      input: "^save/control_dependency"
+      attr {
+        key: "_output_shapes"
+        value {
+          list {
+            shape {
+            }
+          }
+        }
+      }
+    }
+    node {
+      name: "save/RestoreV2/tensor_names"
       op: "Const"
       attr {
         key: "_output_shapes"
         value {
           list {
             shape {
+              dim {
+                size: 1
+              }
             }
           }
         }
@@ -775,6 +1201,9 @@ meta_graphs {
           tensor {
             dtype: DT_STRING
             tensor_shape {
+              dim {
+                size: 1
+              }
             }
             string_val: "a"
           }
@@ -782,13 +1211,16 @@ meta_graphs {
       }
     }
     node {
-      name: "save/restore_slice/shape_and_slice"
+      name: "save/RestoreV2/shape_and_slices"
       op: "Const"
       attr {
         key: "_output_shapes"
         value {
           list {
             shape {
+              dim {
+                size: 1
+              }
             }
           }
         }
@@ -805,6 +1237,9 @@ meta_graphs {
           tensor {
             dtype: DT_STRING
             tensor_shape {
+              dim {
+                size: 1
+              }
             }
             string_val: ""
           }
@@ -812,11 +1247,11 @@ meta_graphs {
       }
     }
     node {
-      name: "save/restore_slice"
-      op: "RestoreSlice"
+      name: "save/RestoreV2"
+      op: "RestoreV2"
       input: "save/Const"
-      input: "save/restore_slice/tensor_name"
-      input: "save/restore_slice/shape_and_slice"
+      input: "save/RestoreV2/tensor_names"
+      input: "save/RestoreV2/shape_and_slices"
       attr {
         key: "_output_shapes"
         value {
@@ -828,15 +1263,11 @@ meta_graphs {
         }
       }
       attr {
-        key: "dt"
-        value {
-          type: DT_FLOAT
-        }
-      }
-      attr {
-        key: "preferred_shard"
+        key: "dtypes"
         value {
-          i: -1
+          list {
+            type: DT_FLOAT
+          }
         }
       }
     }
@@ -844,7 +1275,7 @@ meta_graphs {
       name: "save/Assign"
       op: "Assign"
       input: "a"
-      input: "save/restore_slice"
+      input: "save/RestoreV2"
       attr {
         key: "T"
         value {
@@ -882,13 +1313,16 @@ meta_graphs {
       }
     }
     node {
-      name: "save/restore_slice_1/tensor_name"
+      name: "save/RestoreV2_1/tensor_names"
       op: "Const"
       attr {
         key: "_output_shapes"
         value {
           list {
             shape {
+              dim {
+                size: 1
+              }
             }
           }
         }
@@ -905,6 +1339,9 @@ meta_graphs {
           tensor {
             dtype: DT_STRING
             tensor_shape {
+              dim {
+                size: 1
+              }
             }
             string_val: "b"
           }
@@ -912,13 +1349,16 @@ meta_graphs {
       }
     }
     node {
-      name: "save/restore_slice_1/shape_and_slice"
+      name: "save/RestoreV2_1/shape_and_slices"
       op: "Const"
       attr {
         key: "_output_shapes"
         value {
           list {
             shape {
+              dim {
+                size: 1
+              }
             }
           }
         }
@@ -935,6 +1375,9 @@ meta_graphs {
           tensor {
             dtype: DT_STRING
             tensor_shape {
+              dim {
+                size: 1
+              }
             }
             string_val: ""
           }
@@ -942,11 +1385,11 @@ meta_graphs {
       }
     }
     node {
-      name: "save/restore_slice_1"
-      op: "RestoreSlice"
+      name: "save/RestoreV2_1"
+      op: "RestoreV2"
       input: "save/Const"
-      input: "save/restore_slice_1/tensor_name"
-      input: "save/restore_slice_1/shape_and_slice"
+      input: "save/RestoreV2_1/tensor_names"
+      input: "save/RestoreV2_1/shape_and_slices"
       attr {
         key: "_output_shapes"
         value {
@@ -958,15 +1401,11 @@ meta_graphs {
         }
       }
       attr {
-        key: "dt"
+        key: "dtypes"
         value {
-          type: DT_FLOAT
-        }
-      }
-      attr {
-        key: "preferred_shard"
-        value {
-          i: -1
+          list {
+            type: DT_FLOAT
+          }
         }
       }
     }
@@ -974,7 +1413,7 @@ meta_graphs {
       name: "save/Assign_1"
       op: "Assign"
       input: "b"
-      input: "save/restore_slice_1"
+      input: "save/RestoreV2_1"
       attr {
         key: "T"
         value {
@@ -1012,20 +1451,26 @@ meta_graphs {
       }
     }
     node {
-      name: "save/restore_all"
+      name: "save/restore_shard"
       op: "NoOp"
       input: "^save/Assign"
       input: "^save/Assign_1"
     }
+    node {
+      name: "save/restore_all"
+      op: "NoOp"
+      input: "^save/restore_shard"
+    }
     versions {
       producer: 15
     }
   }
   saver_def {
     filename_tensor_name: "save/Const:0"
-    save_tensor_name: "save/control_dependency:0"
+    save_tensor_name: "save/ShardedFilespec:0"
     restore_op_name: "save/restore_all"
     max_to_keep: 5
+    sharded: true
     keep_checkpoint_every_n_hours: 10000.0
     version: V1
   }
@@ -1048,21 +1493,21 @@ meta_graphs {
     }
   }
   signature_def {
-    key: "regression"
+    key: "tensorflow/serving/regress"
     value {
       inputs {
-        key: "input"
+        key: "inputs"
         value {
-          name: "x:0"
+          name: "tf_example:0"
         }
       }
       outputs {
-        key: "output"
+        key: "outputs"
         value {
-          name: "y:0"
+          name: "Identity:0"
         }
       }
-      method_name: "regression"
+      method_name: "tensorflow/serving/regress"
     }
   }
 }
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 0a87f3306f..d0f0853aa8 100644
--- 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