了解 PyTorch DataLoader Workers 的内存使用情况答案

【问题标题】：Understanding Memory Usage by PyTorch DataLoader Workers了解 PyTorch DataLoader Workers 的内存使用情况
【发布时间】：2020-12-10 20:17:58
【问题描述】：

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

这是否意味着 PyTorch 训练使用 33 个进程 X 15 GB = 495 GB 内存？ htop 显示只有大约 50 GB 的 RAM 和 20 GB 的交换空间正在使用 128 GB RAM 的整台机器上。那么，我们如何解释这种差异呢？

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

谢谢

【问题讨论】：

晚了，但VIRT in htop 大致指的是您的进程可以访问的 RAM 量。而RES 是实际消耗的 RAM。据我了解，RES 是基于父进程的东西——因此请查看父进程的RES 使用情况（将自己设置为树视图）以大致了解您总共使用了多少 RAM。 nvidia-smi 在 GPU 内存方面也是一个很好的代理。
好吧，作者的任何反应都会很棒......

标签： python python-3.x ubuntu deep-learning pytorch

【解决方案1】：

这是否意味着 PyTorch 训练使用 33 个进程 X 15 GB = 495 GB 内存？

没有必要。你有一个工作进程（有几个子进程 - 工作人员）并且 CPU 有几个核心。一名工人通常装载一批。当主进程准备好下一批时，下一批已经被加载并准备好了。这就是加速的秘诀。

我想，你应该使用更少的 num_workers。

了解您的批量大小也会很有趣，您也可以根据训练过程进行调整。

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

我在谷歌上搜索，但找不到具体的公式。我认为这是对你的 CPU 和内存以及 Batch Size 有多少个内核的粗略估计。

要选择 num_workers 取决于您使用的计算机类型、使用的数据集类型以及需要多少动态预处理数据。

HTH

【讨论】：

【解决方案2】：

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

追溯
每个文件名的内存统计数据
计算快照之间的差异

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

【讨论】：