refactor show_traces in memory_tracker (#90145)

zhaojuanmao · pytorchmergebot · commit 743c38554350 · 2023-01-03T22:10:15.000Z
refactor show_tracers in memory_tracker to make it plot multiple figures and also can load serialized stats and then plot figures Pull Request resolved: #90145 Approved by: https://github.com/rohan-varma
diff --git a/torch/distributed/_tools/memory_tracker.py b/torch/distributed/_tools/memory_tracker.py
@@ -146,36 +146,55 @@ def summary(self, top: int = 20) -> None:
         print("------------------------------------------------")
 
     @no_type_check
-    def show_traces(self) -> None:
-        """
-        Show the traces of ``memory_allocated``, ``memory_active`` and ``memory_reserved`` at
-        operator level and the marker 'fw_bw_boundary' at the boundary of forward pass
-        and backward pass.
-        """
+    def show_traces(self, path: str = "") -> None:
+        from itertools import chain
+
         import matplotlib.pyplot as plt
 
-        y_1 = [mb for (name, mb) in self.memories_allocated.values()]
-        y_2 = [mb for (name, mb) in self.memories_active.values()]
-        y_3 = [mb for (name, mb) in self.memories_reserved.values()]
-        min_val = min(y_1 + y_2 + y_3)
-        max_val = max(y_1 + y_2 + y_3)
-        x = list(i for i in range(len(y_1)))
-        fig = plt.figure(figsize=(16, 8))
-        plt.plot(x, list(y_1), label="memory_allocated")
-        plt.plot(x, list(y_2), label="memory_active")
-        plt.plot(x, list(y_3), label="memory_reserved")
-        plt.xlabel("# Operator Calls")
-        plt.ylabel("Memory (MB)")
-        for marker_name, marker in self._markers.items():
-            if marker_name == "fw_bw_boundary":
-                plt.plot(
-                    [marker, marker], [min_val, max_val], "r", lw=2, label=marker_name
-                )
-            else:
-                plt.plot(
-                    [marker, marker], [min_val, max_val], "k-", lw=2, label=marker_name
-                )
-        plt.legend()
+        def _plot_figure(x, y_values, labels):
+            min_val = min(list(chain(*y_values))) * 0.999
+            max_val = max(list(chain(*y_values))) * 1.001
+            plt.figure()
+            for y, label in zip(y_values, labels):
+                plt.plot(x, y, label=label)
+            plt.xlabel("# Operator Calls")
+            plt.ylabel("Memory (MB)")
+            plt.legend()
+            for marker_name, marker in self._markers.items():
+                if marker_name == "fw_bw_boundary":
+                    plt.plot(
+                        [marker, marker],
+                        [min_val, max_val],
+                        "r",
+                        lw=2,
+                        label=marker_name,
+                    )
+                else:
+                    plt.plot(
+                        [marker, marker],
+                        [min_val, max_val],
+                        "k-",
+                        lw=2,
+                        label=marker_name,
+                    )
+
+        if path != "":
+            self.load(path)
+
+        y_1 = [gb for (name, gb) in self.memories_allocated.values()]
+        y_2 = [gb for (name, gb) in self.memories_active.values()]
+        y_3 = [gb for (name, gb) in self.memories_reserved.values()]
+        x = list(range(len(y_1)))
+        # Split figures when there is big difference between
+        # "reserved_memory" and "allocated_memory" or "active_memory".
+        _plot_figure(
+            x,
+            [list(y_1), list(y_2), list(y_3)],
+            ["allocated_memory", "active_memory", "reserved_memory"],
+        )
+        _plot_figure(x, [list(y_1)], ["allocated_memory"])
+        _plot_figure(x, [list(y_2)], ["active_memory"])
+        _plot_figure(x, [list(y_3)], ["reserved_memory"])
 
     def save_stats(self, path: str) -> None:
         """
@@ -190,7 +209,7 @@ def save_stats(self, path: str) -> None:
         }
 
         with open(path, "wb") as f:
-            pickle.dump(stats, f)
+            pickle.dump(stats, f, pickle.HIGHEST_PROTOCOL)
 
     def load(self, path: str) -> None:
         """
