熟悉 AI 训练的人大多知道，在 PyTorch 里，通过 pin_memory 可以加速数据加载和 CPU→GPU 的传输。我在面试里也被问过这个问题：“pin_memory 到底干嘛的？为什么会变快？什么时候该用？”
今天就把这块内容系统地梳理一下。实战中 pin_memory 其实主要有两种用法：

1. 在 DataLoader 里打开 pin_memory=True
2. 手动对张量调用 tensor.pin_memory()

很多坑也恰好就踩在这两种用法上。

内存管理基础

在 CPU 视角下，“内存”其实是虚拟内存 = 物理内存（RAM） + 磁盘交换空间 的抽象：

• 进程看到的是一大片连续的“虚拟地址空间”
• 实际上这些页面（page）可能在 RAM，也可能被换到磁盘上（swap）

这就是所谓的 pageable memory（可分页内存）：

操作系统可以随时把某些页从 RAM 换出到磁盘，再在需要时换回来。

它的好处：

• 让你“看起来”有比实际物理内存大得多的可用空间
• 让 OS 自由调度哪些页在内存、哪些在磁盘

坏处：

• 访问一个当前不在 RAM 里的页会触发 page fault，需要从磁盘换入，延迟巨大
• 对于需要稳定带宽和低延迟的数据传输（比如 GPU DMA），这种不确定性是灾难

与之对应的是 pinned / page-locked memory（固定内存 / 页锁定内存）：

• 这部分内存不会被换出到磁盘，一直“钉死”在物理内存里
• 访问延迟更可预测
• 但数量有限，而且每次 “pin” 的操作本身就比较昂贵

一句话：pinned memory = OS 承诺“这块内存一直在 RAM，不给你换走”。

CUDA 视角

接下来从 GPU 的角度看一眼：CUDA 是怎么把数据从 CPU 搬到显存里的？

• 如果源数据在 pinned 内存：
  GPU 的 DMA 引擎可以直接从这块内存把数据搬到显存，传输路径非常直接。
• 如果源数据在 pageable 内存：
  CUDA 不能直接对 pageable 内存做 DMA，因为 OS 随时可能把这块页换走。
  实际流程是：
  1. CUDA 驱动在内部申请一块临时 pinned buffer
  2. 把 pageable 内存的数据 copy 到这个 buffer
  3. 再从 buffer 通过 DMA 拷贝到 GPU

也就是说，pageable→GPU 实际上是“两次拷贝”：

pageable →（内部 pinned 缓冲区）→ GPU

而 pinned→GPU 是：

pinned → GPU

少了一次拷贝，带宽和延迟自然更好。

所以，

• 直接 pageable_tensor.to('cuda') 比 pinned_tensor.to('cuda') 稍慢
• 因为前者在驱动层偷偷做了一次“隐式 pin + 拷贝”，后者不用。

对张量直接操作

现在从 PyTorch 视角看to()、pin_memory() 和 non_blocking。

to

tensor.to(device)：底层一直是异步拷贝

当你写：

底层用的其实是 cudaMemcpyAsync —— 也就是 异步拷贝。
那为什么感觉它是“同步”的呢？

关键是 PyTorch 在默认情况下帮你做了“拷完就同步”：

• non_blocking=False（默认）：
  • 异步提交 CUDA 拷贝后，紧接着做一次 stream 同步；
  • 对 CPU 来说，这就是一个阻塞调用。
• non_blocking=True：
  • 只提交拷贝，不立刻同步；
  • CPU 线程可以继续向下执行，稍后你自己在合适的地方统一 torch.cuda.synchronize()。

所以重点：

non_blocking=True 不是“让它变成异步”，
而是“别在这一步就等完它”。

pin_memory

PyTorch 提供两种常见方式创建 pinned 张量：

几点要记住：

• pin_memory()会复制一份新张量，不是给原张量“打个标签”；
• 这个复制过程在你调用的那个线程上是 阻塞的：
  • 必须等数据复制到 pinned 内存后，函数才返回。

也就是说：

pin_memory() 本身是一个“挺重”的操作，随手乱用只会拖慢程序。

异步 + pinned 才能真正重叠

要让“传输 + 计算”真正重叠起来，需要同时满足三个条件：

1. 源数据在 pinned memory；
2. 使用的是非默认流或合适的 stream 设计（对多数训练场景 PyTorch 已帮你处理得差不多）；
3. GPU 有空闲 DMA 引擎可用。

对我们平常写训练 loop 来说，核心简化版本就是：

DataLoader(pin_memory=True) + .to(device, non_blocking=True)

一些坑

如果是手动操作的话，还要注意下面 3 点。

1. pin_memory().to() 可能比 to() 还慢

这里做了两件事：

• pin_memory()：在主线程上分配 pinned + 复制一次数据（阻塞）；
• .to("cuda")：再从 pinned 内存异步拷贝到 GPU。

而如果直接 to()：

• CUDA 驱动内部会做一次“隐式 pin + copy”；
• 这套流程是 C++/驱动层高度优化过的，拷贝、缓存、重用策略都比你 Python 层 pin_memory() 要聪明。

所以只用一次的张量，写成 x.pin_memory().to() 多半是在浪费时间。
只有当你反复用同一块 pinned 内存时，这个开销才值得摊平。

2. 一般用 tensor.to(device, non_blocking=True) 就够爽

在多数训练场景中，一种很好的“默认姿势”是：

它的好处是：

• 拷贝本来就是异步的；
• non_blocking=True 让 CPU 可以“先发起多个拷贝，再统一等待”，减少无谓的同步；
• 如果 x、y 本身来自 pinned 内存（比如 DataLoader 帮你 pin 了），传输可以更快、更平滑。

3. CPU→CUDA 可以 non_blocking，CUDA→CPU 小心翻车

两个看起来对称的代码段：

为什么 A 正常、B 却可能出问题？

• 情况 A：
  • .to("cuda", non_blocking=True) 提交 H2D 拷贝；
  • .mean() 这个 kernel 在同一条 CUDA stream 上排队；
  • CUDA 保证同一 stream 上“先拷贝，再计算”，所以数据是完整的。
• 情况 B：
  • .to("cpu", non_blocking=True) 提交 D2H 拷贝后立刻返回一个 CPU 张量；
  • 此时这块 CPU 内存可能还在被填充；
  • .mean() 是 CPU 运算，PyTorch 不会自动给你 torch.cuda.synchronize()；
  • 于是 CPU 直接读了“半成品数据”，结果当然乱。

正确写法应该是：

所以可以记一条小口诀：

CPU→CUDA 的 non_blocking=True 一般安全；
CUDA→CPU 的 non_blocking=True 必须自己同步。

DataLoader 的 pin_memory

“既然 pin_memory() 本身是阻塞操作，那在 DataLoader 里打开 pin_memory=True 为什么还会加速？我们也没在那一步用 non_blocking 啊？”

这里有两个关键点。

1. pin 的阻塞发生在 DataLoader 的 worker 里

当你这样写：

• num_workers > 0 时，DataLoader 会起多个子进程／线程：
  • 负责读数据、做 transform、collate、pin 内存。
• 这些阻塞操作，全都在 worker 里进行，不在你的训练主线程上。

换句话说：

当你的训练 loop 写 for data, target in train_loader 时，
worker 们已经在后台帮你把这个 batch pin 好了。

2. 训练主线程只做：拿 pinned batch → 异步拷贝到 GPU

训练 loop 大概是这样：

时间轴上是：

• GPU 正在训练上一批 batch i；
• DataLoader worker 正在后台准备并 pin 好 batch i+1；
• 当前这批算完后，data.to(device, non_blocking=True) 马上把 batch i+1 往 GPU 上推；
• GPU 几乎无缝衔接下一批，不再干等数据。

默认：

使用之后：

实战建议

1. 这种情况，建议开 pin_memory=True

大致有这么几类：

• 大规模 GPU 训练
  • 大 batch、高分辨率图像／视频、音频等，高吞吐场景；
• 多 GPU / 分布式训练
  • 多张卡同时吃数据，对数据管线延迟非常敏感；
• 低延迟或实时推理
  • 每一毫秒都要抠的时候，pin_memory + non_blocking 能省下不少时间。

搭配方式几乎统一是：

2. 这些情况，可以不用折腾 pin_memory

• 只在 CPU 上训练或推理
  • 开了也没有任何收益，只是多占 RAM；
• 小数据集／轻量任务
  • 数据一下子全塞进 GPU，不存在“GPU 等数据”的瓶颈；
• 机器内存本身就不多（比如 8G）
  • pinned 内存不能被换出，如果你 pin 得太多，系统会非常难受。

3. 手动 pin_memory() 适合的场景

• 实时推理，反复复用固定形状的输入缓冲区；
• 高频 GPU→CPU 回传，自己维护少量 pinned buffer，通过 .copy_() 往里写；
• 自己写异步 pipeline，希望精细控制哪一步 pin、哪一步拷。

不适合的用法：

这基本就属于“在训练主线程上多复制了一遍大 tensor”，收益极低。

总结

pin_memory = 把张量钉死在 RAM 里，让 GPU 传输更顺畅，但他其实是个阻塞操作。
真正好用的姿势是：让 DataLoader 在后台帮你 pin，
然后在训练 loop 里用 .to(device, non_blocking=True) 把传输和计算串成流水线。
少在主线程里乱 pin_memory()，GPU→CPU 的非阻塞传输记得自己同步。