目录

一、num_workers：DataLoader 的“装配线工人数量”

它到底在干什么？

有 / 没有 num_workers，训练流程差在哪？

1️⃣ num_workers=0（最稳，但可能慢）

2️⃣ num_workers > 0（更快，但开始吃资源）

一个非常容易忽略的坑（CT 场景高发）

二、pin_memory=True：给 CPU → GPU 拷贝“铺高速路”

它干什么？

有 / 没有 pin_memory 的流程差异

❌ 不开 pin_memory

✅ 开 pin_memory

但 pinned memory 不是免费的午餐

三、persistent_workers=True：让 worker“全年无休”

它的作用

有 / 没有 persistent_workers 的差异

❌ 关闭（默认）

✅ 开启 persistent_workers

四、三者叠加：为什么你会“训练到一半被 Killed”

五、最务实的训练差异总结（工程视角）

六、稳定优先的推荐起步配置（大规模数据适用）

七、一个很工程师的快速判断法

在 PyTorch 训练里，很多人把注意力都放在模型结构、loss、优化器上，但真正决定你能不能稳定跑完训练的，往往是 DataLoader。

尤其在3D CT / 医学影像 / 大 npz / nii 数据场景中，下面这三个参数一旦组合不当，模型还没学会，进程先被系统一刀Killed：

num_workers pin_memory persistent_workers

一个非常关键的认知先放在最前面：

这三个参数，不改变模型计算本身，只改变“数据如何从磁盘 / CPU → GPU”的流水线。
它们直接影响的是：
吞吐率、CPU/RAM/共享内存占用，以及你会不会被系统干掉。

下面我们把它们当成「训练数据供给系统的三个改装件」，一个一个拆。

一、num_workers：DataLoader 的“装配线工人数量”

它到底在干什么？

num_workers决定 DataLoader 用多少个子进程去执行：

• Dataset.__getitem__

• 数据读取（磁盘 / npz / nii）

• 解码、预处理

• batch 拼装（collate_fn）

简单说：

• num_workers=0
  👉主进程自己干所有数据活

• num_workers>0
  👉 启动 N 个 worker 进程，并行准备 batch

有 / 没有num_workers，训练流程差在哪？

1️⃣num_workers=0（最稳，但可能慢）

流程是完全串行的：

1. 主进程读数据、预处理

2. 拿到一个 batch

3. 拷到 GPU

4. GPU 前向 / 反向

5. 再回到第 1 步

结果是：

• GPU 很可能在等 CPU 喂数据

• 训练吞吐偏低

• 但优点也很明显：

  • 最稳

  • 内存占用最低

  • 基本不碰/dev/shm

  • 调试最友好

2️⃣num_workers > 0（更快，但开始吃资源）

流程变成流水线并行：

• worker 在后台提前准备好 batch（放进队列）

• 主进程直接从队列取 → 马上喂 GPU

效果：

• GPU 等数据的时间明显减少

• 吞吐提升

代价：

• CPU 内存占用上升

• worker 进程数增加（DDP 下是乘法）

• 共享内存/dev/shm压力上来

一个非常容易忽略的坑（CT 场景高发）

PyTorch 默认：

prefetch_factor = 2

也就是说：

每个 worker 会提前准备 2 个 batch

内存占用近似是：

num_workers × prefetch_factor × batch_size × 单样本大小

再叠加：

• 3D CT（单样本动辄几十 MB）

• Dataset 里缓存了大对象

• DDP / accelerate

👉 内存会涨得比你想象快得多。

二、pin_memory=True：给 CPU → GPU 拷贝“铺高速路”

它干什么？

当你把数据从 CPU 拷贝到 GPU（.to(device)）时：

• 普通内存：拷贝慢一点

• Pinned（页锁定）内存：拷贝更快，可和 GPU 计算重叠

pin_memory=True的作用是：

DataLoader 先把 batch 放进 pinned memory，再从 pinned 拷到 GPU

有 / 没有 pin_memory 的流程差异

❌ 不开 pin_memory

1. worker 产出普通内存 batch

2. 主进程同步拷贝到 GPU

3. GPU 计算

✅ 开 pin_memory

1. worker 产出 batch

2. DataLoader额外复制一份到 pinned 内存

3. 主进程从 pinned 内存拷到 GPU（更快）

如果你在训练代码里这样写：

batch = batch.to(device, non_blocking=True)

效果会更明显（前提是 pinned）。

但 pinned memory 不是免费的午餐

它的问题在于：

• CPU 侧多了一次拷贝

• pinned 内存是锁页内存，系统不容易回收

• 一旦 pinned 累积过多：

  👉更容易 OOM 或被系统 Killed

在以下组合下尤其危险：

• num_workers > 0

• 有预取（prefetch）

• 单样本很大

• 长时间训练

三、persistent_workers=True：让 worker“全年无休”

它的作用

默认情况下：

• 每个 epoch 结束

• DataLoader 会关闭 worker

• 下个 epoch 再重新 fork / spawn

开启：

persistent_workers=True

意味着：

worker 在整个训练期间都不退出

有 / 没有 persistent_workers 的差异

❌ 关闭（默认）

每个 epoch：

• worker 退出

• 资源释放

• 下个 epoch 重新启动

优点：

• 内存更容易被回收

• Dataset 里的问题不会长期累积

• 更稳

缺点：

• epoch 切换有开销

✅ 开启 persistent_workers

优点：

• epoch 切换更快

• 数据供给更平滑

但代价非常隐蔽，也最致命：

worker 内的任何“泄漏”，都会一直累积

包括但不限于：

• Dataset 里缓存的对象

• 未释放的 numpy buffer

• np.load 的隐式缓存

• pinned memory 的滞留

• 预取队列堆积

结果就是：

前面跑得好好的，跑到一半突然被 Killed
没有 traceback，没有报错，只有冷冰冰的一行字。

四、三者叠加：为什么你会“训练到一半被 Killed”

你当前的典型组合是：

• num_workers = 2

• pin_memory = True

• persistent_workers = True

• prefetch_factor = 2（默认）

这会制造一种非常经典的曲线：

CPU 内存 / SHM 使用量缓慢上升 → 突破阈值 → SIGKILL

尤其在这些条件下：

• 3D 医学影像

• batch_size 不小

• DDP / accelerate

• Docker / 云平台/dev/shm很小

这不是模型“突然抽风”，
这是数据管道在悄悄囤积资源，直到系统忍无可忍。

五、最务实的训练差异总结（工程视角）

一句话版本：

• num_workers=0：最稳，但可能慢

• num_workers>0：更快，但吃内存

• pin_memory=True：GPU 训练更快，但 CPU 压力上升

• persistent_workers=True：短期爽，长期风险最大

六、稳定优先的推荐起步配置（大规模数据适用）

如果你现在的目标是：先稳定跑通训练，建议从这里开始：

train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=1, pin_memory=True, persistent_workers=False, # 先关 prefetch_factor=1, # 非常关键 drop_last=True, )

调优顺序建议是：

1. 先确认这套能稳定跑完整个训练

2. 再把num_workers：1 → 2

3. 最后才考虑persistent_workers=True

4. 一旦“跑一半死”：

   • 第一时间降prefetch_factor

   • 或直接关 persistent_workers

七、一个很工程师的快速判断法

记住这条就够了：

如果是“跑一半才被杀”，优先怀疑：
persistent_workers + 预取队列 + pin_memory 的内存/SHM 累积问题。

模型真出问题，一般会给你 traceback。
系统杀进程，才是这种一句话不说的冷处理。