Add more display options to benchmark, including FLOPs

Change: 143266630

1 files changed, 27 insertions, 0 deletions

diff --git a/tensorflow/core/util/stat_summarizer.h b/tensorflow/core/util/stat_summarizer.h
index 7159cce88d..4568fa420b 100644
--- a/tensorflow/core/util/stat_summarizer.h
+++ b/tensorflow/core/util/stat_summarizer.h
@@ -108,6 +108,29 @@ class Stat {
   HighPrecisionValueType squared_sum_ = 0;
 };
 
+// Used to control the output of the statistics summarizer;
+class StatSummarizerOptions {
+ public:
+  StatSummarizerOptions()
+      : show_run_order(true),
+        run_order_limit(0),
+        show_time(true),
+        time_limit(10),
+        show_memory(true),
+        memory_limit(10),
+        show_type(true),
+        show_summary(true) {}
+
+  bool show_run_order;
+  int run_order_limit;
+  bool show_time;
+  int time_limit;
+  bool show_memory;
+  int memory_limit;
+  bool show_type;
+  bool show_summary;
+};
+
 // A class intended to make performance analysis easier by collecting StepStats
 // and showing in an easily understandable format where CPU time is being spent.
 // See tensorflow/examples/android/jni/tensorflow_jni.cc for an example usage.
@@ -124,6 +147,9 @@ class StatSummarizer {
 
   explicit StatSummarizer(const tensorflow::GraphDef& tensorflow_graph);
 
+  StatSummarizer(const tensorflow::GraphDef& tensorflow_graph,
+                 const StatSummarizerOptions& options);
+
   // Adds another run's StepStats output to the aggregate counts.
   void ProcessStepStats(const StepStats& step_stats);
 
@@ -184,6 +210,7 @@ class StatSummarizer {
   std::vector<string> nodes_in_def_order_;
   std::map<std::string, Detail> details_;
   std::map<string, string> node_types_;
+  StatSummarizerOptions options_;
 };
 
 }  // namespace tensorflow