如何获取 pytorch 在 CPU/主内存上的内存统计信息？

Question

有时你需要知道你的程序在高峰期间需要多少内存，但可能不太关心这个高峰何时发生以及持续多长时间等。Pytorch 有一个很好的工具，可以报告在 GPU 上运行时的内存使用情况，您只需在程序结束时调用一次：

memory_usage = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
torch.cuda.reset_peak_memory_stats()

这段代码非常简单，因为它使您无需运行一个单独的线程每毫秒监视您的内存并找到峰值。

现在我的问题是：为什么这只适用于 GPU？我找不到像 torch.cpu.memory_stats() 这样的东西。在CPU上运行时，这个挂件是什么？

Answer 1

为此，您需要使用 Pytorch Profiler，它可以为您提供有关 CPU 和内存消耗的详细信息。

更多细节：

https://pytorch.org/blog/introducing-pytorch-profiler-the-new-and-improved-performance-tool/

https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html