pytorch：参数pin_memory=True和non_blocking=True的作用

最新推荐文章于 2025-11-08 10:40:53 发布

原创最新推荐文章于 2025-11-08 10:40:53 发布 · 1.2w 阅读

48 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #人工智能

pytorch 专栏收录该内容

32 篇文章

订阅专栏

本文介绍如何通过设置pin_memory和non_blocking参数来提高PyTorch数据加载和传输效率，实现训练过程加速。pin_memory用于配置数据是否存放于锁页内存以保持与GPU的高速数据交换，non_blocking则确保数据在GPU内存中不被释放，两者结合使用可以显著提升模型训练的速度。

部署运行你感兴趣的模型镜像

一、pin_memory

pin_memory是dataloader()的参数，默认值为False，其作用是是否把把数据存放在锁页内存中。主机的内存根据物理内存（内存条）与虚拟内存（硬盘）进行数据交换分为锁页内存和不锁页内存：

锁页内存：数据存放在物理内存上（内存条）上；
不锁页内存：当物理内存（内存条）满载时，把部分数据转换到虚拟内存上（硬盘）上。
锁页内存（pin_memory)能够保持与GPU进行高速传输，在训练时加快数据的读取，从而加快训练速度。因此，如果主机/服务器的内存足够大，建议把pin_memory设为True，如：

trainloader = torch.utils.data.DataLoader(dataset=traindata, batch_size=BATCH_SIZE, shuffle=True, num_workers=1, pin_memory=True)

二、non_blocking

non_blocking时cuda()的参数，默认值为False，其作用和pin_memory一样，pin_memory是针对物理内存（内存条），而non_blocking是针对GPU上的内存（显存），表士把数据锁页在显存上，在后台进程过程中不释放。一般地，如果pin_momery为True，把non_blocking也设为True，有助于加速数据传输，加快训练过程，如：