Updated benchmark program to handle multiple inputs and outputs

Change: 138471362

3 files changed, 103 insertions, 73 deletions

diff --git a/tensorflow/tools/benchmark/benchmark_model.cc b/tensorflow/tools/benchmark/benchmark_model.cc
index 3279967aaa..c3d9e865b1 100644
--- a/tensorflow/tools/benchmark/benchmark_model.cc
+++ b/tensorflow/tools/benchmark/benchmark_model.cc
@@ -78,36 +78,35 @@ Status InitializeSession(int num_threads, const string& graph,
   return Status::OK();
 }
 
-Status RunBenchmark(DataType input_data_type, TensorShape input_shape,
-                    const string& input_layer, const string output_layer,
-                    Session* session, StatSummarizer* stats) {
-  Tensor input_tensor(input_data_type, input_shape);
-
-  switch (input_data_type) {
-    case DT_INT32: {
-      auto int_tensor = input_tensor.flat<int32>();
-      int_tensor = int_tensor.constant(0.0);
-      break;
+Status RunBenchmark(const std::vector<InputLayerInfo>& inputs,
+                    const std::vector<string>& outputs, Session* session,
+                    StatSummarizer* stats) {
+  std::vector<std::pair<string, tensorflow::Tensor> > input_tensors;
+  for (const InputLayerInfo& input : inputs) {
+    Tensor input_tensor(input.data_type, input.shape);
+    switch (input.data_type) {
+      case DT_INT32: {
+        auto int_tensor = input_tensor.flat<int32>();
+        int_tensor = int_tensor.constant(0.0);
+        break;
+      }
+      case DT_FLOAT: {
+        auto float_tensor = input_tensor.flat<float>();
+        float_tensor = float_tensor.constant(0.0);
+        break;
+      }
+      case DT_QUINT8: {
+        auto int_tensor = input_tensor.flat<quint8>();
+        int_tensor = int_tensor.constant(0.0);
+        break;
+      }
+      default:
+        LOG(FATAL) << "Unsupported input type: " << input.data_type;
     }
-    case DT_FLOAT: {
-      auto float_tensor = input_tensor.flat<float>();
-      float_tensor = float_tensor.constant(0.0);
-      break;
-    }
-    case DT_QUINT8: {
-      auto int_tensor = input_tensor.flat<quint8>();
-      int_tensor = int_tensor.constant(0.0);
-      break;
-    }
-    default:
-      LOG(FATAL) << "Unsupported input type: " << input_data_type;
+    input_tensors.push_back({input.name, input_tensor});
   }
 
-  std::vector<std::pair<string, tensorflow::Tensor> > input_tensors(
-      {{input_layer, input_tensor}});
-
   std::vector<tensorflow::Tensor> output_tensors;
-  std::vector<string> output_names({output_layer});
 
   tensorflow::Status s;
 
@@ -115,8 +114,8 @@ Status RunBenchmark(DataType input_data_type, TensorShape input_shape,
   run_options.set_trace_level(RunOptions::FULL_TRACE);
   RunMetadata run_metadata;
 
-  s = session->Run(run_options, input_tensors, output_names, {},
-                   &output_tensors, &run_metadata);
+  s = session->Run(run_options, input_tensors, outputs, {}, &output_tensors,
+                   &run_metadata);
 
   if (!s.ok()) {
     LOG(ERROR) << "Error during inference: " << s;
@@ -132,9 +131,9 @@ Status RunBenchmark(DataType input_data_type, TensorShape input_shape,
 }
 
 Status TimeMultipleRuns(double sleep_seconds, int num_runs,
-                        DataType input_data_type, TensorShape input_shape,
-                        const string& input_layer, const string output_layer,
-                        Session* session, StatSummarizer* stats) {
+                        const std::vector<InputLayerInfo>& inputs,
+                        const std::vector<string>& outputs, Session* session,
+                        StatSummarizer* stats) {
   // Convert the run_delay string into a timespec.
   timespec req;
   req.tv_sec = static_cast<time_t>(sleep_seconds);
@@ -142,8 +141,7 @@ Status TimeMultipleRuns(double sleep_seconds, int num_runs,
 
   LOG(INFO) << "Running benchmark";
   for (int i = 0; i < num_runs; ++i) {
-    Status run_status = RunBenchmark(input_data_type, input_shape, input_layer,
-                                     output_layer, session, stats);
+    Status run_status = RunBenchmark(inputs, outputs, session, stats);
     if (!run_status.ok()) {
       LOG(INFO) << "Failed on run " << i;
       return run_status;
@@ -162,10 +160,10 @@ Status TimeMultipleRuns(double sleep_seconds, int num_runs,
 
 int Main(int argc, char** argv) {
   string graph = "/data/local/tmp/tensorflow_inception_graph.pb";
-  string input_layer = "input:0";
-  string input_layer_shape = "1,224,224,3";
-  string input_layer_type = "float";
-  string output_layer = "output:0";
+  string input_layer_string = "input:0";
+  string input_layer_shape_string = "1,224,224,3";
+  string input_layer_type_string = "float";
+  string output_layer_string = "output:0";
   int num_runs = 50;
   string run_delay = "-1.0";
   int num_threads = -1;
@@ -175,10 +173,10 @@ int Main(int argc, char** argv) {
 
   std::vector<Flag> flag_list = {
       Flag("graph", &graph, "graph file name"),
-      Flag("input_layer", &input_layer, "input layer name"),
-      Flag("input_layer_shape", &input_layer_shape, "input layer shape"),
-      Flag("input_layer_type", &input_layer_type, "input layer type"),
-      Flag("output_layer", &output_layer, "output layer name"),
+      Flag("input_layer", &input_layer_string, "input layer names"),
+      Flag("input_layer_shape", &input_layer_shape_string, "input layer shape"),
+      Flag("input_layer_type", &input_layer_type_string, "input layer type"),
+      Flag("output_layer", &output_layer_string, "output layer name"),
       Flag("num_runs", &num_runs, "number of runs"),
       Flag("run_delay", &run_delay, "delay between runs in seconds"),
       Flag("num_threads", &num_threads, "number of threads"),
@@ -194,6 +192,28 @@ int Main(int argc, char** argv) {
     return -1;
   }
 
+  std::vector<string> input_layers = str_util::Split(input_layer_string, ',');
+  std::vector<string> input_layer_shapes =
+      str_util::Split(input_layer_shape_string, ':');
+  std::vector<string> input_layer_types =
+      str_util::Split(input_layer_type_string, ',');
+  std::vector<string> output_layers = str_util::Split(output_layer_string, ',');
+  if ((input_layers.size() != input_layer_shapes.size()) ||
+      (input_layers.size() != input_layer_types.size())) {
+    LOG(ERROR) << "There must be the same number of items in --input_layer,"
+               << " --input_layer_shape, and --input_layer_type, for example"
+               << " --input_layer=input1,input2 --input_layer_type=float,float "
+               << " --input_layer_shape=1,224,224,4:1,20";
+    LOG(ERROR) << "--input_layer=" << input_layer_string << " ("
+               << input_layers.size() << " items)";
+    LOG(ERROR) << "--input_layer_type=" << input_layer_type_string << " ("
+               << input_layer_types.size() << " items)";
+    LOG(ERROR) << "--input_layer_shape=" << input_layer_shape_string << " ("
+               << input_layer_shapes.size() << " items)";
+    return -1;
+  }
+  const size_t inputs_count = input_layers.size();
+
   ::tensorflow::port::InitMain(argv[0], &argc, &argv);
   if (argc > 1) {
     LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage;
@@ -201,10 +221,10 @@ int Main(int argc, char** argv) {
   }
 
   LOG(INFO) << "Graph: [" << graph << "]";
-  LOG(INFO) << "Input layer: [" << input_layer << "]";
-  LOG(INFO) << "Input shape: [" << input_layer_shape << "]";
-  LOG(INFO) << "Input type: [" << input_layer_type << "]";
-  LOG(INFO) << "Output layer: [" << output_layer << "]";
+  LOG(INFO) << "Input layers: [" << input_layer_string << "]";
+  LOG(INFO) << "Input shapes: [" << input_layer_shape_string << "]";
+  LOG(INFO) << "Input types: [" << input_layer_type_string << "]";
+  LOG(INFO) << "Output layers: [" << output_layer_string << "]";
   LOG(INFO) << "Num runs: [" << num_runs << "]";
   LOG(INFO) << "Inter-run delay (seconds): [" << run_delay << "]";
   LOG(INFO) << "Num threads: [" << num_threads << "]";
@@ -221,21 +241,26 @@ int Main(int argc, char** argv) {
   }
 
   const double sleep_seconds = std::strtod(run_delay.c_str(), nullptr);
-  DataType input_data_type;
-  CHECK(DataTypeFromString(input_layer_type, &input_data_type))
-      << input_layer_type << " was an invalid type";
-  std::vector<int32> sizes;
-  CHECK(str_util::SplitAndParseAsInts(input_layer_shape, ',', &sizes))
-      << "Incorrect size string specified: " << input_layer_shape;
-  TensorShape input_shape;
-  for (int i = 0; i < sizes.size(); ++i) {
-    input_shape.AddDim(sizes[i]);
+
+  std::vector<InputLayerInfo> inputs;
+  for (int n = 0; n < inputs_count; ++n) {
+    InputLayerInfo input;
+    CHECK(DataTypeFromString(input_layer_types[n], &input.data_type))
+        << input_layer_types[n] << " was an invalid type";
+    std::vector<int32> sizes;
+    CHECK(str_util::SplitAndParseAsInts(input_layer_shapes[n], ',', &sizes))
+        << "Incorrect size string specified: " << input_layer_shapes[n];
+    for (int i = 0; i < sizes.size(); ++i) {
+      input.shape.AddDim(sizes[i]);
+    }
+    input.name = input_layers[n];
+    inputs.push_back(input);
   }
 
   const int64 start_time = Env::Default()->NowMicros();
   Status time_status =
-      TimeMultipleRuns(sleep_seconds, num_runs, input_data_type, input_shape,
-                       input_layer, output_layer, session.get(), stats.get());
+      TimeMultipleRuns(sleep_seconds, num_runs, inputs, output_layers,
+                       session.get(), stats.get());
   const int64 end_time = Env::Default()->NowMicros();
   const double wall_time = (end_time - start_time) / 1000000.0;
 
@@ -252,13 +277,10 @@ int Main(int argc, char** argv) {
 
   if (!benchmark_name.empty() && !output_prefix.empty()) {
     // Compute the total number of values per input.
-    int64 total_size = 1;
-    for (int32 size : sizes) {
-      total_size *= size;
-    }
+    int64 total_size = inputs[0].shape.num_elements();
 
     // Throughput in MB/s
-    const double throughput = DataTypeSize(input_data_type) * total_size *
+    const double throughput = DataTypeSize(inputs[0].data_type) * total_size *
                               num_runs / static_cast<double>(wall_time) /
                               (1024 * 1024);
 
diff --git a/tensorflow/tools/benchmark/benchmark_model.h b/tensorflow/tools/benchmark/benchmark_model.h
index c7de398547..e09e4e2606 100644
--- a/tensorflow/tools/benchmark/benchmark_model.h
+++ b/tensorflow/tools/benchmark/benchmark_model.h
@@ -22,21 +22,28 @@ limitations under the License.
 namespace tensorflow {
 namespace benchmark_model {
 
+// Used to help construct dummy inputs for the benchmarking.
+struct InputLayerInfo {
+  string name;
+  DataType data_type;
+  TensorShape shape;
+};
+
 // Loads a model from disk into a new session, and sets up the stats collection.
 Status InitializeSession(int num_threads, const string& graph,
                          std::unique_ptr<Session>* session,
                          std::unique_ptr<StatSummarizer>* stats);
 
 // Does a single run of the model that's been loaded into the given session.
-Status RunBenchmark(DataType input_data_type, TensorShape input_shape,
-                    const string& input_layer, const string output_layer,
-                    Session* session, StatSummarizer* stats);
+Status RunBenchmark(const std::vector<InputLayerInfo>& inputs,
+                    const std::vector<string>& outputs, Session* session,
+                    StatSummarizer* stats);
 
 // Runs the model multiple time, keeping track of timing information.
 Status TimeMultipleRuns(double sleep_seconds, int num_runs,
-                        DataType input_data_type, TensorShape input_shape,
-                        const string& input_layer, const string output_layer,
-                        Session* session, StatSummarizer* stats);
+                        const std::vector<InputLayerInfo>& inputs,
+                        const std::vector<string>& outputs, Session* session,
+                        StatSummarizer* stats);
 
 // Handles all setup and argument parsing.
 int Main(int argc, char** argv);
diff --git a/tensorflow/tools/benchmark/benchmark_model_test.cc b/tensorflow/tools/benchmark/benchmark_model_test.cc
index cbbec7bbbe..f74d79ee43 100644
--- a/tensorflow/tools/benchmark/benchmark_model_test.cc
+++ b/tensorflow/tools/benchmark/benchmark_model_test.cc
@@ -33,7 +33,9 @@ TEST(BenchmarkModelTest, InitializeAndRun) {
   // Create a simple graph and write it to filename_pb.
   const int input_width = 400;
   const int input_height = 10;
-  const TensorShape input_shape({input_width, input_height});
+  benchmark_model::InputLayerInfo input;
+  input.shape = TensorShape({input_width, input_height});
+  input.data_type = DT_FLOAT;
   const TensorShape constant_shape({input_height, input_width});
 
   Tensor constant_tensor(DT_FLOAT, constant_shape);
@@ -41,8 +43,8 @@ TEST(BenchmarkModelTest, InitializeAndRun) {
 
   auto root = Scope::NewRootScope().ExitOnError();
   auto placeholder =
-      ops::Placeholder(root, DT_FLOAT, ops::Placeholder::Shape(input_shape));
-  const string input_name = placeholder.node()->name();
+      ops::Placeholder(root, DT_FLOAT, ops::Placeholder::Shape(input.shape));
+  input.name = placeholder.node()->name();
   auto m = ops::MatMul(root, placeholder, constant_tensor);
   const string output_name = m.node()->name();
 
@@ -58,9 +60,8 @@ TEST(BenchmarkModelTest, InitializeAndRun) {
   TF_ASSERT_OK(
       benchmark_model::InitializeSession(1, filename_pb, &session, &stats));
 
-  TF_ASSERT_OK(benchmark_model::TimeMultipleRuns(0.0, 10, DT_FLOAT, input_shape,
-                                                 input_name, output_name,
-                                                 session.get(), stats.get()));
+  TF_ASSERT_OK(benchmark_model::TimeMultipleRuns(
+      0.0, 10, {input}, {output_name}, session.get(), stats.get()));
 }
 
 }  // namespace