Add ExtractExampleParserConfiguration method

- Extract the fixed length and variable length feature configurations
  output tensor names from a given GraphDef.
- This will allow for the use case of bypassing an unnecessary tensorflow.Example
  serialize/deserialize at serving/inference time by extracting the configuration,
  running  the proto -> tensor helpers directly and feeding the graph with
  the properly named tensors
Change: 122636456

4 files changed, 689 insertions, 0 deletions

diff --git a/tensorflow/core/example/example_parser_configuration.cc b/tensorflow/core/example/example_parser_configuration.cc
new file mode 100644
index 0000000000..798ebfd364
--- /dev/null
+++ b/tensorflow/core/example/example_parser_configuration.cc
@@ -0,0 +1,160 @@
+/* Copyright 2016 Google Inc. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/example/example_parser_configuration.h"
+
+#include <vector>
+
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/example/feature.pb_text.h"
+#include "tensorflow/core/framework/numeric_op.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/protobuf.h"
+
+namespace tensorflow {
+
+Status FindNodeIndexByName(const tensorflow::GraphDef& graph,
+                           const string& node_name, int* node_idx) {
+  for (int i = 0; i < graph.node_size(); ++i) {
+    const auto& node = graph.node(i);
+    if (node.name() == node_name) {
+      *node_idx = i;
+      return Status::OK();
+    }
+  }
+  return errors::InvalidArgument(node_name, " not found in GraphDef");
+}
+
+Status ExtractExampleParserConfiguration(
+    const tensorflow::GraphDef& graph, const string& node_name,
+    tensorflow::Session* session,
+    std::vector<FixedLenFeature>* fixed_len_features,
+    std::vector<VarLenFeature>* var_len_features) {
+  int node_idx;
+  TF_RETURN_IF_ERROR(FindNodeIndexByName(graph, node_name, &node_idx));
+
+  const auto& node = graph.node(node_idx);
+  if (node.op() != "ParseExample") {
+    return errors::InvalidArgument(node_name, " node is not a ParseExample op");
+  }
+
+  auto& attr_map = node.attr();
+  auto num_sparse = attr_map.at("Nsparse").i();
+  auto num_dense = attr_map.at("Ndense").i();
+  fixed_len_features->resize(num_dense);
+  var_len_features->resize(num_sparse);
+
+  auto tdense = attr_map.at("Tdense");
+  auto dense_shapes = attr_map.at("dense_shapes");
+  auto sparse_types = attr_map.at("sparse_types");
+
+  // Consistency check attributes.
+  if (tdense.list().type_size() != num_dense) {
+    return errors::InvalidArgument("Node attr Tdense has ",
+                                   tdense.list().type_size(),
+                                   " elements != Ndense attr: ", num_dense);
+  }
+
+  if (dense_shapes.list().shape_size() != num_dense) {
+    return errors::InvalidArgument("Node attr dense_shapes has ",
+                                   dense_shapes.list().shape_size(),
+                                   " elements != Ndense attr: ", num_dense);
+  }
+
+  if (sparse_types.list().type_size() != num_sparse) {
+    return errors::InvalidArgument("Node attr sparse_types has ",
+                                   sparse_types.list().type_size(),
+                                   " elements != NSparse attr: ", num_sparse);
+  }
+
+  for (int i = 0; i < tdense.list().type_size(); ++i) {
+    (*fixed_len_features)[i].dtype = tdense.list().type(i);
+    // Convert TensorShapeProto to TensorShape.
+    (*fixed_len_features)[i].shape = TensorShape(dense_shapes.list().shape(i));
+  }
+
+  for (int i = 0; i < sparse_types.list().type_size(); ++i) {
+    (*var_len_features)[i].dtype = sparse_types.list().type(i);
+  }
+
+  // We must fetch the configuration input tensors to the ParseExample op.
+  // Skipping index = 0, which is the serialized proto input.
+  std::vector<string> fetch_names(node.input_size() - 1);
+  for (int i = 1; i < node.input_size(); ++i) {
+    fetch_names[i - 1] = node.input(i);
+  }
+
+  std::vector<Tensor> op_input_tensors;
+
+  TF_RETURN_IF_ERROR(session->Run({},               // no_inputs,
+                                  fetch_names, {},  // no target_node_names,
+                                  &op_input_tensors));
+
+  // The input tensors are laid out sequentially in a flat manner.
+  // Here are the various start offsets.
+  int sparse_keys_start = 1;
+  int dense_keys_start = sparse_keys_start + num_sparse;
+  int dense_defaults_start = dense_keys_start + num_dense;
+
+  for (int i = 0; i < num_sparse; ++i) {
+    int input_idx = sparse_keys_start + i;
+    (*var_len_features)[i].key = op_input_tensors[input_idx].scalar<string>()();
+  }
+
+  for (int i = 0; i < num_dense; ++i) {
+    FixedLenFeature& config = (*fixed_len_features)[i];
+    int dense_keys_offset = dense_keys_start + i;
+    config.key = op_input_tensors[dense_keys_offset].scalar<string>()();
+
+    int defaults_offset = dense_defaults_start + i;
+    config.default_value = op_input_tensors[defaults_offset];
+  }
+
+  // The output tensors are laid out sequentially in a flat manner.
+  // Here are the various start offsets.
+  int sparse_indices_output_start = 0;
+  int sparse_values_output_start = sparse_indices_output_start + num_sparse;
+  int sparse_shapes_output_start = sparse_values_output_start + num_sparse;
+  int dense_values_output_start = sparse_shapes_output_start + num_sparse;
+
+  string node_output_prefix = strings::StrCat(node_name, ":");
+
+  for (int i = 0; i < num_sparse; ++i) {
+    VarLenFeature& config = (*var_len_features)[i];
+
+    int indices_offset = sparse_indices_output_start + i;
+    config.indices_output_tensor_name =
+        strings::StrCat(node_output_prefix, indices_offset);
+
+    int values_offset = sparse_values_output_start + i;
+    config.values_output_tensor_name =
+        strings::StrCat(node_output_prefix, values_offset);
+
+    int shapes_offset = sparse_shapes_output_start + i;
+    config.shapes_output_tensor_name =
+        strings::StrCat(node_output_prefix, shapes_offset);
+  }
+
+  for (int i = 0; i < num_dense; ++i) {
+    int output_idx = dense_values_output_start + i;
+    (*fixed_len_features)[i].values_output_tensor_name =
+        strings::StrCat(node_output_prefix, output_idx);
+  }
+  return Status::OK();
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/core/example/example_parser_configuration.h b/tensorflow/core/example/example_parser_configuration.h
new file mode 100644
index 0000000000..cf90dd7f4b
--- /dev/null
+++ b/tensorflow/core/example/example_parser_configuration.h
@@ -0,0 +1,47 @@
+/* Copyright 2016 Google Inc. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef THIRD_PARTY_TENSORFLOW_CORE_EXAMPLE_EXAMPLE_PARSER_CONFIGURATION_H_
+#define THIRD_PARTY_TENSORFLOW_CORE_EXAMPLE_EXAMPLE_PARSER_CONFIGURATION_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/util/example_proto_helper.h"
+#include "tensorflow/core/util/sparse/sparse_tensor.h"
+
+// This is a set of helper methods that will make it possible to share
+// tensorflow::Example proto Tensor conversion code inside the ExampleParserOp
+// OpKernel as well as in external code.
+namespace tensorflow {
+
+// Given a graph and the node_name of a ParseExample op,
+// extract the FixedLenFeature/VarLenFeature configurations.
+Status ExtractExampleParserConfiguration(
+    const tensorflow::GraphDef& graph, const string& node_name,
+    tensorflow::Session* session,
+    std::vector<FixedLenFeature>* fixed_len_features,
+    std::vector<VarLenFeature>* var_len_features);
+
+}  // namespace tensorflow
+
+#endif  // THIRD_PARTY_TENSORFLOW_CORE_EXAMPLE_EXAMPLE_PARSE_CONFIGURATION_H_
diff --git a/tensorflow/core/example/example_parser_configuration_test.cc b/tensorflow/core/example/example_parser_configuration_test.cc
new file mode 100644
index 0000000000..e804b123f2
--- /dev/null
+++ b/tensorflow/core/example/example_parser_configuration_test.cc
@@ -0,0 +1,148 @@
+/* Copyright 2016 Google Inc. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/example/example_parser_configuration.h"
+
+#include "tensorflow/core/example/example.pb.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/util/example_proto_helper.h"
+
+namespace tensorflow {
+namespace {
+
+void ReadFileToStringOrDie(Env* env, const string& filename, string* output) {
+  TF_CHECK_OK(ReadFileToString(env, filename, output));
+}
+Session* CreateSession() {
+  SessionOptions options;
+  (*options.config.mutable_device_count())["CPU"] = 2;
+  return NewSession(options);
+}
+
+class ExtractExampleParserConfigurationTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    string proto_string;
+    string filename =
+        io::JoinPath(testing::TensorFlowSrcRoot(),
+                     "core/example/testdata/parse_example_graph_def.pbtxt");
+    ReadFileToStringOrDie(Env::Default(), filename, &proto_string);
+    protobuf::TextFormat::ParseFromString(proto_string, &graph_def_);
+    session_.reset(CreateSession());
+    session_->Create(graph_def_);
+  }
+
+  NodeDef* parse_example_node() {
+    for (int i = 0; i < graph_def_.node_size(); ++i) {
+      auto mutable_node = graph_def_.mutable_node(i);
+      if (mutable_node->name() == "ParseExample/ParseExample") {
+        return mutable_node;
+      }
+    }
+    return nullptr;
+  }
+  GraphDef graph_def_;
+  std::unique_ptr<Session> session_;
+};
+
+TEST_F(ExtractExampleParserConfigurationTest, OpNotFound) {
+  std::vector<FixedLenFeature> dense_vec;
+  std::vector<VarLenFeature> sparse_vec;
+  Status status = ExtractExampleParserConfiguration(
+      graph_def_, "BlarseExample/ParseExample", session_.get(), &dense_vec,
+      &sparse_vec);
+
+  EXPECT_EQ(status.code(), error::INVALID_ARGUMENT);
+}
+
+TEST_F(ExtractExampleParserConfigurationTest, InconsistentAttrNsparse) {
+  std::vector<FixedLenFeature> dense_vec;
+  std::vector<VarLenFeature> sparse_vec;
+
+  NodeDef* node = parse_example_node();
+  auto mutable_attr = node->mutable_attr();
+  (*mutable_attr)["Nsparse"].set_i(3);
+
+  Status status = ExtractExampleParserConfiguration(
+      graph_def_, "ParseExample/ParseExample", session_.get(), &dense_vec,
+      &sparse_vec);
+
+  EXPECT_EQ(status.code(), error::INVALID_ARGUMENT);
+}
+
+TEST_F(ExtractExampleParserConfigurationTest, InconsistentAttrNdense) {
+  std::vector<FixedLenFeature> dense_vec;
+  std::vector<VarLenFeature> sparse_vec;
+
+  NodeDef* node = parse_example_node();
+  auto mutable_attr = node->mutable_attr();
+  (*mutable_attr)["Ndense"].set_i(2);
+
+  Status status = ExtractExampleParserConfiguration(
+      graph_def_, "ParseExample/ParseExample", session_.get(), &dense_vec,
+      &sparse_vec);
+
+  EXPECT_EQ(status.code(), error::INVALID_ARGUMENT);
+}
+
+TEST_F(ExtractExampleParserConfigurationTest, Basic) {
+  std::vector<FixedLenFeature> dense_vec;
+  std::vector<VarLenFeature> sparse_vec;
+  Status status = ExtractExampleParserConfiguration(
+      graph_def_, "ParseExample/ParseExample", session_.get(), &dense_vec,
+      &sparse_vec);
+
+  EXPECT_EQ(Status::OK(), status);
+  EXPECT_EQ(2, sparse_vec.size());
+  EXPECT_EQ(3, dense_vec.size());
+
+  EXPECT_EQ("sf0", sparse_vec[0].key);
+  EXPECT_EQ(DT_STRING, sparse_vec[0].dtype);
+  EXPECT_EQ("ParseExample/ParseExample:0",
+            sparse_vec[0].indices_output_tensor_name);
+  EXPECT_EQ("ParseExample/ParseExample:2",
+            sparse_vec[0].values_output_tensor_name);
+  EXPECT_EQ("ParseExample/ParseExample:4",
+            sparse_vec[0].shapes_output_tensor_name);
+
+  EXPECT_EQ("sf1", sparse_vec[1].key);
+  EXPECT_EQ(DT_STRING, sparse_vec[1].dtype);
+  EXPECT_EQ("ParseExample/ParseExample:1",
+            sparse_vec[1].indices_output_tensor_name);
+  EXPECT_EQ("ParseExample/ParseExample:3",
+            sparse_vec[1].values_output_tensor_name);
+  EXPECT_EQ("ParseExample/ParseExample:5",
+            sparse_vec[1].shapes_output_tensor_name);
+
+  EXPECT_EQ("x", dense_vec[0].key);
+  EXPECT_EQ(DT_FLOAT, dense_vec[0].dtype);
+  EXPECT_EQ("ParseExample/ParseExample:6",
+            dense_vec[0].values_output_tensor_name);
+
+  EXPECT_EQ("y", dense_vec[1].key);
+  EXPECT_EQ(DT_FLOAT, dense_vec[1].dtype);
+  EXPECT_EQ("ParseExample/ParseExample:7",
+            dense_vec[1].values_output_tensor_name);
+
+  EXPECT_EQ("z", dense_vec[2].key);
+  EXPECT_EQ(DT_FLOAT, dense_vec[2].dtype);
+  EXPECT_EQ("ParseExample/ParseExample:8",
+            dense_vec[2].values_output_tensor_name);
+}
+
+}  // namespace
+}  // namespace tensorflow
diff --git a/tensorflow/core/example/testdata/parse_example_graph_def.pbtxt b/tensorflow/core/example/testdata/parse_example_graph_def.pbtxt
new file mode 100644
index 0000000000..437499e73d
--- /dev/null
+++ b/tensorflow/core/example/testdata/parse_example_graph_def.pbtxt
@@ -0,0 +1,334 @@
+node {
+  name: "batch_constant"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_STRING
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+        }
+        string_val: "z"
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/Reshape/shape"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_FLOAT
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 1
+          }
+        }
+        float_val: 1
+      }
+    }
+  }
+}
+
+node {
+  name: "ParseExample/key_y"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_FLOAT
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+        }
+        float_val: 0.0
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/Reshape_1/shape"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_FLOAT
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 1
+          }
+        }
+        float_val: 1
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/key_z"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_FLOAT
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+        }
+        float_val: 0.0
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/Reshape_2/shape"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_FLOAT
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 1
+          }
+        }
+        float_val: 1
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/ParseExample/names"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_STRING
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+          dim {
+          }
+        }
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/ParseExample/sparse_keys_0"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_STRING
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+        }
+        string_val: "sf0"
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/ParseExample/sparse_keys_1"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_STRING
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+        }
+        string_val: "sf1"
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/ParseExample/dense_keys_0"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_STRING
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+        }
+        string_val: "x"
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/ParseExample/dense_keys_1"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_STRING
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+        }
+        string_val: "y"
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/ParseExample/dense_keys_2"
+  op: "Const"
+  attr {
+    key: "dtype"
+    value {
+      type: DT_STRING
+    }
+  }
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_STRING
+        tensor_shape {
+        }
+        string_val: "z"
+      }
+    }
+  }
+}
+node {
+  name: "ParseExample/ParseExample"
+  op: "ParseExample"
+  input: "batch_constant"
+  input: "ParseExample/ParseExample/names"
+  input: "ParseExample/ParseExample/sparse_keys_0"
+  input: "ParseExample/ParseExample/sparse_keys_1"
+  input: "ParseExample/ParseExample/dense_keys_0"
+  input: "ParseExample/ParseExample/dense_keys_1"
+  input: "ParseExample/ParseExample/dense_keys_2"
+  input: "ParseExample/Reshape/shape"
+  input: "ParseExample/Reshape_1/shape"
+  input: "ParseExample/Reshape_2/shape"
+  attr {
+    key: "Nsparse"
+    value {
+      i: 2
+    }
+  }
+  attr {
+    key: "Ndense"
+    value {
+      i: 3
+    }
+  }
+  attr {
+    key: "sparse_types"
+    value {
+      list {
+        type: DT_STRING
+        type: DT_STRING
+      }
+    }
+  }
+  attr {
+    key: "Tdense"
+    value {
+      list {
+        type: DT_FLOAT
+        type: DT_FLOAT
+        type: DT_FLOAT
+      }
+    }
+  }
+  attr {
+    key: "dense_shapes"
+    value {
+      list {
+        shape {
+          dim {
+            size: 1
+          }
+        }
+        shape {
+          dim {
+            size: 1
+          }
+        }
+        shape {
+          dim {
+            size: 1
+          }
+        }
+      }
+    }
+  }
+}
+versions {
+  producer: 9
+}
+