Add python wrapper for StatSummarizer.

Change: 138933733

1 files changed, 12 insertions, 6 deletions

diff --git a/tensorflow/core/util/stat_summarizer.cc b/tensorflow/core/util/stat_summarizer.cc
index 33cf5fe961..0b675eaac9 100644
--- a/tensorflow/core/util/stat_summarizer.cc
+++ b/tensorflow/core/util/stat_summarizer.cc
@@ -338,13 +338,19 @@ std::string StatSummarizer::GetStatsByOrderOfNodeDefinitions(
   return stream.str();
 }
 
+std::string StatSummarizer::GetOutputString() const {
+  std::stringstream stream;
+  stream << "Total time (us): " << run_total_micros_;
+  stream << GetTimingStatsByRunOrder();
+  stream << GetTimingStatsByTopDurations();
+  stream << "Total Memory (bytes): " << memory_;
+  stream << GetMemoryStatsByRunOrder();
+  stream << GetMemoryStatsByUsage();
+  return stream.str();
+}
+
 void StatSummarizer::PrintStepStats() const {
-  LOG(INFO) << "Total time (us): " << run_total_micros_;
-  LOG(INFO) << GetTimingStatsByRunOrder();
-  LOG(INFO) << GetTimingStatsByTopDurations();
-  LOG(INFO) << "Total Memory (bytes): " << memory_;
-  LOG(INFO) << GetMemoryStatsByRunOrder();
-  LOG(INFO) << GetMemoryStatsByUsage();
+  LOG(INFO) << GetOutputString();
   LOG(INFO);
 }