数据加载优化：DataLoader的num_workers、pin_memory与prefetch_factor

数据加载，说白了就是给GPU喂饭。你模型再强，算力再猛，如果数据送不上去，GPU就得干等着。我见过太多同学，模型结构调得飞起，结果训练速度被数据加载卡死——GPU利用率只有30%，你说气不气人？

今天咱们就聊聊PyTorch DataLoader里三个最关键的参数：num_workers、pin_memory和prefetch_factor。调好了，训练速度翻倍不是梦。

1. num_workers：到底开几个进程合适？

num_workers控制的是数据加载的子进程数量。默认是0，意味着主进程自己干所有活——这通常是最慢的。

我个人习惯，先看机器有多少CPU核心。比如你有个8核的CPU，GPU是3090，那num_workers设成4或6是比较稳妥的。为什么不是8？因为操作系统、PyTorch本身、数据预处理都要占用CPU资源，全占满反而会互相打架。

我的经验： 从 num_workers=4 开始试，观察GPU利用率。如果GPU利用率一直低于80%，就往上加。加到某个点发现GPU利用率不再提升，甚至下降了——那就是瓶颈到了。

我曾经在一个项目里，把num_workers从4调到8，结果训练反而变慢了。排查了半天，发现是磁盘IO扛不住了。多个进程同时读硬盘，磁盘队列满了，每个请求都在排队。嗯，这里要注意：num_workers不是越大越好，它受限于CPU、内存带宽和磁盘IO。

一个简单的经验公式：

num_workers = 4 * GPU数量   # 单GPU场景
num_workers = 2 * GPU数量   # 多GPU场景（避免进程间通信开销）

但这只是起点，最终还是要靠实验来调。

2. pin_memory：锁页内存到底锁了什么？

pin_memory=True，说白了就是把数据固定在物理内存里，不让操作系统把它交换到磁盘上。为什么这么做？因为GPU要从CPU内存拷贝数据时，如果数据在虚拟内存里（可能被换到磁盘），那拷贝速度会慢得离谱。

你想想看，GPU要数据，结果发现数据在磁盘上，得先读回内存，再拷贝到显存——这中间多了一道IO，速度直接掉一个数量级。

  核心结论： 只要你的内存够用，pin_memory=True 一定要开。它几乎不增加额外开销，但能显著提升数据从CPU到GPU的传输速度。

我记得有一次帮同事调一个图像分类模型，他死活不开pin_memory，说「内存不够」。我一看，他机器有64GB内存，数据集才8GB。我说你这不是内存不够，是舍不得用。开了之后，训练速度提升了15%。

什么时候不开？只有一种情况：你的内存非常紧张，连锁页内存都分配不出来。但说实话，现在做深度学习，32GB内存是标配，64GB很常见，基本不用担心这个问题。

3. prefetch_factor：预取多少数据才够？

prefetch_factor是PyTorch 1.7之后才暴露出来的参数。它控制每个worker预取多少个batch的数据。默认是2，意味着每个worker会提前准备2个batch的数据放在队列里。

为什么要预取？因为数据加载和模型训练是并行的。模型在算当前batch的时候，DataLoader已经在准备下一个batch了。如果预取不够，模型算完了还得等数据——这就是所谓的「数据饥饿」。

我个人的调优思路：

默认值2：大多数场景下够用，尤其是数据预处理不复杂的时候
调大到4或8：如果你的数据预处理很重（比如实时数据增强、复杂的transform），预取多一点可以掩盖预处理的时间
调小到1：内存非常紧张的时候，或者你的数据加载已经很快了，不需要那么多预取

注意： prefetch_factor 调大会增加内存占用。每个worker预取的数据都放在内存里，如果batch size很大（比如256或512），预取8个batch可能吃掉好几个GB的内存。我曾经在一个项目里把prefetch_factor设成16，结果内存爆了，训练直接OOM。

一个更直观的理解：

# 默认配置
DataLoader(dataset, batch_size=64, num_workers=4, prefetch_factor=2)
# 每个worker预取2个batch，总共预取 4 * 2 = 8 个batch
# 内存占用 ≈ 8 * 64 * 每个样本大小

# 激进配置
DataLoader(dataset, batch_size=64, num_workers=4, prefetch_factor=8)
# 每个worker预取8个batch，总共预取 4 * 8 = 32 个batch
# 内存占用 ≈ 32 * 64 * 每个样本大小

4. 三个参数的协同调优

这三个参数不是孤立的，它们互相影响。我总结了一个调优流程：

先定num_workers：从4开始，观察GPU利用率，逐步增加直到瓶颈
开启pin_memory：只要内存够，就设成True
调整prefetch_factor：从2开始，如果GPU利用率还是上不去，尝试4或8

下面这张图展示了数据加载的完整流程：

5. 实战配置示例

下面是我常用的几个配置模板，你可以直接拿来用：

# 场景1：小数据集，单GPU，内存充足
DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,
    pin_memory=True,
    prefetch_factor=2
)

# 场景2：大数据集，多GPU，内存紧张
DataLoader(
    dataset,
    batch_size=128,
    num_workers=8,
    pin_memory=True,
    prefetch_factor=1  # 内存不够，少预取
)

# 场景3：数据增强很重，需要掩盖预处理时间
DataLoader(
    dataset,
    batch_size=64,
    num_workers=6,
    pin_memory=True,
    prefetch_factor=4  # 多预取，让worker提前干活
)

一个小技巧： 如果你用 torch.utils.data.DataLoader 的 persistent_workers=True，可以避免每个epoch重新创建worker进程，能省下不少时间。配合 prefetch_factor 一起用效果更好。

最后说一句，数据加载优化这件事，没有银弹。不同的硬件、不同的数据集、不同的预处理流程，最优配置都不一样。我的建议是：先跑起来，再调优。别一开始就纠结参数，先让模型能跑，然后用nvidia-smi看GPU利用率，用htop看CPU和内存，一步步调。

我曾经在一个项目里，花了两天时间调数据加载，把训练速度从200 it/s提到了350 it/s。虽然过程有点折腾，但看到GPU利用率从40%飙到95%的时候，那种成就感，嗯，值得。

  一句话总结： num_workers管并行度，pin_memory管传输速度，prefetch_factor管预取深度。三个参数配合好，你的GPU就不会再「饿肚子」了。

📚 课程目录

数据加载优化：DataLoader的num_workers、pin_memory与prefetch_factor

1. num_workers：到底开几个进程合适？

2. pin_memory：锁页内存到底锁了什么？

3. prefetch_factor：预取多少数据才够？

4. 三个参数的协同调优

5. 实战配置示例