理解PyTorch DataLoader工作进程的内存使用

Question

当使用 num_workers=32 运行 PyTorch 训练程序时，htop 显示出 33 个 Python 进程，每个进程都有 32 GB 的 VIRT 和 15 GB 的 RES。

这是否意味着 PyTorch 训练正在使用 33 个进程 X 15 GB = 495 GB 的内存？但是，htop 仅显示整个计算机上使用了约 50 GB 的 RAM 和 20 GB 的 swap，而该机器有 128 GB 的 RAM。那么，我们如何解释这种差异？

是否有更准确的方法来计算主 PyTorch 程序及其所有子 DataLoader 工作进程使用的总内存量？

谢谢

- Athena Wisdom

---

晚了一些，但是htop中的VIRT大致指的是您的进程可以访问的RAM数量。而RES则是实际消耗的RAM。据我所知，RES是基于父进程的，因此查看父进程的RES使用情况（将自己设置为树形视图）可以大致了解您总共使用了多少RAM。在GPU内存方面，nvidia-smi也是一个很好的代理。 - John

希望作者能够回复任何反馈。 - j35t3r

Answer 1

这是否意味着PyTorch训练使用了33个进程X 15 GB = 495 GB的内存？ 不一定。您有一个工作进程（带有几个子进程-工作人员），CPU具有多个核心。通常，一个工作进程加载一个批次。当主进程准备好另一个批次时，下一个批次可能已经准备好了。这是加速的秘诀。 我猜，您应该使用更少的num_workers。 也很有趣知道您的批量大小，您也可以为培训过程进行调整。 有没有更精确的计算主PyTorch程序及其所有子DataLoader工作程序使用的总RAM量的方法？ 我在Google上搜索，但找不到具体的公式。我认为这是您的CPU和内存以及批量大小有多少个核心的粗略估计。 选择num_workers取决于您使用的计算机类型，采用的数据集类型以及您的数据需要多少即时预处理。 HTH

- j35t3r

Answer 2

有一个名为tracemalloc的Python函数，用于跟踪分配给Python的内存块。https://docs.python.org/3/library/tracemalloc.html

• Tracebacks
• 每个文件的内存静态信息
• 计算快照之间的差异

import tracemalloc
tracemalloc.start()
do_someting_that_consumes_ram_and releases_some()
# show how much RAM the above code allocated and the peak usage
current, peak =  tracemalloc.get_traced_memory()
print(f"{current:0.2f}, {peak:0.2f}")
tracemalloc.stop()

https://discuss.pytorch.org/t/measuring-peak-memory-usage-tracemalloc-for-pytorch/34067

- CoffeeBaconAddict