Show output tensor sizes in model benchmark tool

Change: 129708903

1 files changed, 6 insertions, 2 deletions

diff --git a/tensorflow/core/util/stat_summarizer.h b/tensorflow/core/util/stat_summarizer.h
index c5dea66c65..d0130db2d4 100644
--- a/tensorflow/core/util/stat_summarizer.h
+++ b/tensorflow/core/util/stat_summarizer.h
@@ -115,6 +115,9 @@ class StatSummarizer {
   // can be pasted into a spreadsheet for further analysis.
   void PrintStepStats() const;
 
+  // Prints the output tensor sizes and types for each node.
+  void PrintOutputs() const;
+
   // Summarizes all nodes' stat in the order of node names defined in the graph.
   std::string GetStatsByOrderOfNodeDefinitions() const;
 
@@ -130,7 +133,7 @@ class StatSummarizer {
 
   void Reset() {
     run_total_micros_.Reset();
-    timing_details_.clear();
+    details_.clear();
   }
 
   // Returns number of runs.
@@ -144,6 +147,7 @@ class StatSummarizer {
     int64 first_start_micros;
     int64 first_rel_end_micros;
     int64 total_micros;
+    std::vector<TensorDescription> outputs;
   };
 
   enum struct SortingMetric {
@@ -163,7 +167,7 @@ class StatSummarizer {
   int64 first_node_start_micros_;
   Stat<int64> run_total_micros_;
   std::vector<string> nodes_in_def_order_;
-  std::map<std::string, Detail> timing_details_;
+  std::map<std::string, Detail> details_;
   std::map<string, string> node_types_;
 };