了解 PyTorch DataLoader Worker 的内存使用情况

Question

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

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

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

谢谢

Answer 1

python 有一个名为tracemalloc 的函数，用于跟踪分配给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