Pytorch 中 pin_memory 和 non_blocking 用法

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

我们知道,pin_memorynon_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 是指“不要阻塞”,希望运算能够同时进行。那么什么样的运算能够同时进行呢?——比如数

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值