Tweak debugger result (apache#4426)

zxy844288792 · Dec 13, 2019 · 153cc94 · 153cc94
1 parent b4e406c
commit 153cc94
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/python/tvm/contrib/debugger/debug_result.py b/python/tvm/contrib/debugger/debug_result.py
@@ -204,9 +204,8 @@ def dump_graph_json(self, graph):
         with open(os.path.join(self._dump_path, graph_dump_file_name), 'w') as outfile:
             json.dump(graph, outfile, indent=4, sort_keys=False)
 
-    def display_debug_result(self, sort_by_time=True):
-        """Displays the debugger result"
-        """
+    def get_debug_result(self, sort_by_time=True):
+        """Return the debugger result"""
         header = ["Node Name", "Ops", "Time(us)", "Time(%)", "Shape", "Inputs", "Outputs"]
         lines = ["---------", "---", "--------", "-------", "-----", "------", "-------"]
         eid = 0
@@ -244,10 +243,15 @@ def display_debug_result(self, sort_by_time=True):
                 if item_len > max_len:
                     max_len = item_len
             fmt = fmt + "{:<" + str(max_len + 2) + "}"
-        print(fmt.format(*header))
-        print(fmt.format(*lines))
+        log = [fmt.format(*header)]
+        log.append(fmt.format(*lines))
         for row in data:
-            print(fmt.format(*row))
+            log.append(fmt.format(*row))
+        return '\n'.join(log)
+
+    def display_debug_result(self, sort_by_time=True):
+        """Displays the debugger result"""
+        print(self.get_debug_result(sort_by_time))
 
 def save_tensors(params):
     """Save parameter dictionary to binary bytes.