[Profiler] Support HTML plot output for profiler export_memory_timeli…

…ne API (pytorch#99751)

Summary:
Pull Request resolved: pytorch#99751

Support the file extension .html, which will include a PNG image of the plot embedded into an HTML file.

This allows users to avoid processing the timeline manually in their own frontend UI.

Test Plan:
CI Tests

Ran on resnet50 model and generated this html file w/ plot:
See attached html file: {F954232276}
Screenshot: {F954232469}

Reviewed By: davidberard98

Differential Revision: D45152735

Pulled By: aaronenyeshi

fbshipit-source-id: 6b4f00678ee804b093391f5ba8150a526a4eed9c

torch/profiler/_memory_profiler.py

@@ -1015,3 +1015,51 @@ def export_memory_timeline(self, path, device) -> None:
         import json
         with open(path, 'w') as f:
             json.dump([times, sizes], f)
+
+    def export_memory_timeline_html(self, path, device, figsize=(20, 12), title=None) -> None:
+        """Exports the memory timeline as an HTML file which contains
+        the memory timeline plot embedded as a PNG file."""
+        # Check if user has matplotlib installed, return gracefully if not.
+        import importlib.util
+        matplotlib_spec = importlib.util.find_spec("matplotlib")
+        if matplotlib_spec is None:
+            print("export_memory_timeline_html failed because matplotlib was not found.")
+            return
+
+        import matplotlib.pyplot as plt
+        import numpy as np
+        from base64 import b64encode
+        from tempfile import NamedTemporaryFile
+        from os import remove
+
+        mt = self._coalesce_timeline(device)
+        times, sizes = np.array(mt[0]), np.array(mt[1])
+        stacked = np.cumsum(sizes, axis=1) / 1024**3
+
+        # Plot memory timeline as stacked data
+        fig = plt.figure(figsize=figsize, dpi=80)
+        axes = fig.gca()
+        for category, color in _CATEGORY_TO_COLORS.items():
+            i = _CATEGORY_TO_INDEX[category]
+            axes.fill_between(
+                times / 1e3, stacked[:, i], stacked[:, i + 1], color=color, alpha=0.7
+            )
+        fig.legend(["Unknown" if i is None else i.name for i in _CATEGORY_TO_COLORS])
+        axes.set_xlabel("Time (us)")
+        axes.set_ylabel("Memory (GB)")
+        title = "\n\n".join(
+            ([title] if title else []) + [f"Max: {stacked[:, -1].max():.2f} GB"]
+        )
+        axes.set_title(title)
+
+        # Embed the memory timeline image into the HTML file
+        tmpfile = NamedTemporaryFile('wb', suffix='.png', delete=False)
+        tmpfile.close()
+        fig.savefig(tmpfile.name, format='png')
+
+        with open(tmpfile.name, 'rb') as tmp:
+            encoded = b64encode(tmp.read()).decode('utf-8')
+            html = 'GPU Memory Timeline HTML' + '<img src=\'data:image/png;base64,{}\'>'.format(encoded)
+            with open(path, 'w') as f:
+                f.write(html)
+        remove(tmpfile.name)

torch/profiler/profiler.py

@@ -245,7 +245,10 @@ def export_memory_timeline(self, path: str, device: str = None) -> None:
         self.mem_tl = MemoryProfileTimeline(self._memory_profile())
 
         # Depending on the file suffix, save the data as json.gz or json.
-        if path.endswith('.gz'):
+        # For html, we can embed the image into an HTML file.
+        if path.endswith('.html'):
+            self.mem_tl.export_memory_timeline_html(path, device)
+        elif path.endswith('.gz'):
             fp = tempfile.NamedTemporaryFile('w+t', suffix='.json', delete=False)
             fp.close()
             self.mem_tl.export_memory_timeline(fp.name, device)