I’m wondering if there is a way to view CPU memory timelines, similar to what is done in this blog post for CUDA memory.
I tried using the torch.profiler’s export_memory_timeline function, with device=cpu, but it seems that this function is deprecated, and it didn’t work with html format (with json format it seemed to work but then I don’t know how to visualize the result).
Ideally, I would like a result similar to that (copied from the blog post linked above):
The PyTorch profiler can profile memory, but as far as I know, it doesn’t offer any tool to visualize it.
As a workaround, I made a script to extract the “[memory]” events from some saved traces (exported by profiler.export_chrome_trace) and to plot them in a matplotlib plot. It’s not interactive, and it’s hard to tell from which operation the memory allocation came from, but at least it helped me in my use-case.
Here is the how the resulting plot looks like (here, I’m comparing two different implementations of a method in terms of memory usage over time):
