我们知道,pin_memory 和 non_blocking 可以帮助加速 Pytorch 训练过程。这篇文章里,我以一个比较粗浅的理解分析一下,它们为什么能够加速训练。
pin_memory = True
当我们要在 GPU 上进行训练时,自然需要把数据从CPU(一般情况下,数据存储在 CPU 上)转移到 GPU 上。但是 CPU 与 GPU 之间的数据交互是比较慢的,特别是 CPU 中的 pageable memory (可分页内存)与 GPU 之间的交互,这个过程需要建立一个临时缓冲区(pinned memory)——如下图左边所示。
在 dataloader 中 设置 pin_memory=True,程序一开始就把数据放在 pinned memory (锁页内存)上,这样向 GPU 传输数据时就会更快,如上图右边所示。这样一来,就避免了 CPU 内部内存的拷贝,节约了时间。
注意: pin_memory=True 虽好,也不要“贪杯”噢。
If you overuse pinned memory, it can cause serious problems when running low on RAM, and you should be aware that pinning is often an expensive operation. ——Use pinned memory buffers
一般设备间传输速度的数量级:
锁页内存和 GPU 显存之间的拷贝速度大约是 6GB/s
可分页内存和 GPU 显存间的拷贝速度大约是 3GB/s。
GPU 内存间速度是 30GB/s,CPU 间内存速度是 10GB/s ——详解 Pytorch 里的 pin_memory 和 non_blocking
non_blocking = True
顾名思义,non_blocking 是指“不要阻塞”,希望运算能够同时进行。那么什么样的运算能够同时进行呢?——比如数

文章介绍了Pytorch中pin_memory和non_blocking两个设置如何帮助加速训练过程。pin_memory通过将数据加载到锁页内存以优化GPU与CPU间的数据传输速度,而non_blocking则允许数据传输与计算操作并行进行,提高效率。配合dataloader的num_workers设置,可以进一步提升数据加载速度。


被折叠的 条评论
为什么被折叠?



