3、数据加载优化：使用DataLoader的num_workers、pin_memory、prefetch_factor

数据加载，往往是训练流程中最容易被忽视的瓶颈。

很多人把精力都花在调模型、改架构上，结果GPU利用率只有30%——说白了，显卡在等数据吃饭。我早期做项目时就踩过这个坑，模型写得再花哨，数据喂不进去，一切都是白搭。

这一节，我们就来聊聊PyTorch DataLoader里的三个关键参数：num_workers、pin_memory、prefetch_factor。它们能帮你把数据加载的延迟降到最低，让GPU真正跑起来。

3.1 为什么数据加载会成为瓶颈？

训练一个深度学习模型，流程大致是这样的：

  1. CPU从硬盘/内存读取数据
  2. CPU对数据进行预处理（裁剪、归一化、增强等）
  3. 数据从CPU拷贝到GPU
  4. GPU开始计算

问题出在哪？GPU算得快，但CPU读数据慢。尤其是当你的数据集很大、预处理逻辑复杂时，CPU就成了那个拖后腿的。

你想想看，GPU一秒钟能算几百张图，但CPU一秒钟只能读几十张——那GPU剩下的时间就只能干等着。这就是所谓的「I/O瓶颈」。

核心思路：让数据加载和模型计算「流水线化」。GPU在算当前batch的时候，CPU提前把下一个batch准备好。这样GPU就不用等了。

DataLoader的三个参数，就是用来实现这个流水线的。

3.2 num_workers：用多进程加速数据读取

num_workers 是最直观的参数。它控制着用几个子进程来加载数据。

默认值是0，意思是只用主进程加载。这其实是最慢的模式——主进程既要管数据，又要管模型，分身乏术。

# 单进程加载（慢）
dataloader = DataLoader(dataset, batch_size=32, num_workers=0)

# 多进程加载（快）
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

那到底设多少合适？

我个人习惯是：num_workers 设为CPU核心数的一半或全部。比如你的机器有8核，设4或8都可以。

但要注意，不是越大越好。我在项目中遇到过，把num_workers设到16，结果CPU被占满，反而拖慢了训练。因为进程切换也是有开销的。

CPU核心数 推荐 num_workers 说明
4 2 ~ 4 小数据集够用
8 4 ~ 8 大多数场景的甜点区
16 8 ~ 12 留一些核心给系统和其他进程
32+ 16 ~ 24 大规模训练时可以考虑

小技巧：如果你用Windows系统，num_workers 设大于0时可能会报错。这是因为Windows的进程启动方式和Linux不同。解决办法：把数据加载逻辑放到 if __name__ == '__main__': 里面。

3.3 pin_memory：加速CPU到GPU的数据传输

pin_memory 这个参数，很多人不理解它的作用。我简单解释一下。

CPU的内存分两种：可分页内存（pageable）和锁页内存（pinned）。默认情况下，PyTorch张量存在可分页内存里。当你要把数据从CPU拷贝到GPU时，系统得先把它搬到锁页内存，再拷贝过去——多了一步操作。

而 pin_memory=True 的意思是：直接在锁页内存里创建张量。这样拷贝到GPU时就省掉了中间步骤，速度能快不少。

# 开启锁页内存
dataloader = DataLoader(dataset, batch_size=32, pin_memory=True)

什么时候该用？

  • 你的GPU显存够大，训练batch size不小
  • 你的数据预处理不复杂，瓶颈主要在数据传输上
  • 你用的是NVIDIA GPU（CUDA环境）

注意：锁页内存是有限的。如果你把 pin_memory=True 和很大的 num_workers 一起用，可能会把锁页内存占满，导致程序崩溃。我曾经在32核服务器上设了 num_workers=32 加 pin_memory=True，结果系统直接OOM了。后来降到16才稳定。

3.4 prefetch_factor：预取多少个batch

prefetch_factor 是PyTorch 1.7之后加入的参数。它控制每个worker预取多少个batch的数据。

默认值是2。意思是每个worker会提前加载2个batch的数据放在队列里。这样当GPU需要下一个batch时，数据已经准备好了。

# 默认预取2个batch
dataloader = DataLoader(dataset, batch_size=32, num_workers=4, prefetch_factor=2)

# 预取更多batch（适合数据加载特别慢的场景）
dataloader = DataLoader(dataset, batch_size=32, num_workers=4, prefetch_factor=4)

调大 prefetch_factor 的好处：

  • 数据加载更平滑，不容易出现「断流」
  • 适合预处理逻辑复杂的场景

坏处：

  • 占用更多内存（每个worker多存几个batch）
  • 如果数据增强是随机的，预取太多可能导致数据分布偏差（不过影响很小）

我个人建议：如果你的数据加载时间占训练时间的比例超过20%，可以试试把 prefetch_factor 从2调到4或8。但别调太大，否则内存扛不住。

3.5 三个参数如何搭配？

这三个参数不是孤立的。它们一起决定了数据加载管线的效率。

下面这张图展示了它们之间的关系：

DataLoader 数据加载流水线 硬盘/存储 读取 Worker 1 Worker 2 Worker 3 Worker 4 num_workers=4 预处理 预取队列 prefetch_factor=2 拷贝 GPU pin_memory=True 流水线工作流程： 1. 多个Worker并行从硬盘读取数据（num_workers控制并行度） 2. 预处理后的数据放入预取队列（prefetch_factor控制队列深度） 3. 数据通过锁页内存快速拷贝到GPU（pin_memory加速传输）

嗯，这里要注意：这三个参数的设置没有绝对标准。我一般会先设 num_workers=4、pin_memory=True、prefetch_factor=2，然后观察GPU利用率。如果利用率低于80%，就逐步调大 num_workers 和 prefetch_factor。

3.6 实战配置示例

下面是一个完整的DataLoader配置，适合大多数中等规模的训练任务：

from torch.utils.data import DataLoader

# 推荐配置（8核CPU，中等规模数据集）
dataloader = DataLoader(
    dataset=my_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,          # CPU核心数的一半
    pin_memory=True,        # 开启锁页内存
    prefetch_factor=2,      # 每个worker预取2个batch
    persistent_workers=True # 训练完一个epoch不销毁worker进程
)

补充说明：persistent_workers=True 是另一个有用的参数。它让worker进程在epoch之间保持存活，避免反复创建销毁的开销。配合 num_workers 一起用效果更好。

3.7 避坑指南

我曾经在调优时遇到过几个坑，分享给你：

  • 坑1：num_workers 设太大导致内存爆炸。 每个worker都会复制一份数据集和预处理逻辑。如果数据集本身很大，worker太多会把内存撑爆。解决办法：用 torch.utils.data.IterableDataset 或者减少worker数。
  • 坑2：pin_memory 和 num_workers 一起用时，锁页内存不够。 每个worker都会占用锁页内存。如果总锁页内存超过系统限制，程序会卡死。我一般会先算一下：每个batch大小 × prefetch_factor × num_workers，确保这个值不超过系统锁页内存的80%。
  • 坑3：数据增强太复杂，worker处理不过来。 如果每个worker处理一个batch需要很长时间，那设再多worker也没用。这时候应该优化预处理逻辑，或者用 torchvision.transforms 的GPU版本。

数据加载优化，说白了就是让CPU和GPU各司其职、并行工作。调好这三个参数，你的训练速度能提升30%到50%，甚至更多。别小看这一步，很多时候模型训练慢，不是模型的问题，是数据没喂进去。

← 上一章 📖 返回目录 下一章 →