2. 数据加载优化

数据加载，说白了就是给GPU喂饭。喂得太慢，GPU就得饿着等，算力白白浪费。我见过太多项目，模型写得花里胡哨，结果一跑起来，GPU利用率才30%，瓶颈全在数据加载上。今天咱们就把这块彻底聊透。

DataLoader的核心参数调优

PyTorch的DataLoader，我几乎每个项目都在用。但很多人只是无脑设个batch_size就完事了。其实它的参数门道不少。

batch_size：这个不用多说，直接影响显存占用。我习惯先设一个较小的值（比如8或16），跑一个batch看看显存占用，再逐步往上加。注意，不是越大越好——batch太大，模型收敛反而变慢。

num_workers：这是多进程加载的核心参数。它控制着同时有几个子进程在预处理数据。我个人的经验是：

• CPU核心数的一半到三分之二，是个不错的起点
• 如果CPU资源充足，可以设为核心数-1（留一个给系统）
• 设得太多反而会因进程切换开销导致性能下降

pin_memory机制：让数据传输更快

pin_memory=True，这个参数我建议你养成习惯加上。为什么？

正常情况下，CPU内存中的数据是分页的，GPU要读取时，得先拷贝到固定内存区域，再传输到GPU。这个过程有额外开销。而pin_memory=True，相当于告诉操作系统：这块内存别动它，固定住。这样GPU可以直接通过DMA（直接内存访问）读取，速度提升明显。

我曾经在一个图像分类项目里做过对比：

看到没？就改一个参数，性能提升超过30%。而且这几乎不消耗额外资源，属于白捡的优化。

prefetch_factor：预取多少才合适

prefetch_factor这个参数，很多人可能都没注意过。它控制着每个worker预先加载多少个batch的数据。

默认值是2，意味着每个worker会预取2个batch。如果你的数据预处理很耗时（比如图像增强、文本tokenization），适当增大这个值能让GPU一直有活干。

我建议的调优思路：

• 先设成2，跑一个epoch看看GPU利用率
• 如果GPU经常空闲，逐步增加到4、8、16
• 注意内存占用——prefetch_factor越大，占用的内存越多

举个例子，我有一个NLP项目，每个样本都要做复杂的tokenization。默认prefetch_factor=2时，GPU利用率只有60%。我改成8之后，利用率直接飙到95%。代价是内存从8GB涨到了14GB——但值得。

多进程加载策略

多进程加载，说白了就是让多个CPU核心同时干活。但这里有个坑：不是所有场景都适合多进程。

什么时候用多进程？

• 数据预处理是CPU密集型（如图像resize、数据增强）
• 数据量很大，单进程加载跟不上GPU速度
• 你有足够的CPU核心和内存

什么时候别用？

• 数据预处理很简单（比如只是读个numpy数组）
• 数据量很小，加载不是瓶颈
• 内存紧张（多进程会复制数据，内存占用翻倍）

我曾经在一个视频处理项目里踩过坑。视频帧的预处理特别耗时，我一开始设了8个worker，结果内存直接爆了——因为每个worker都在缓存视频帧。后来我改成4个worker，配合prefetch_factor=1，才稳定下来。

实战配置模板

下面是我常用的DataLoader配置模板，你可以直接拿去用：

from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,           # CPU核心数的一半
    pin_memory=True,         # 固定内存，加速GPU传输
    prefetch_factor=4,       # 每个worker预取4个batch
    persistent_workers=True, # 训练完一个epoch不销毁worker
    drop_last=True           # 丢弃最后一个不完整的batch
)

这里有个参数我特别想提一下：persistent_workers=True。默认情况下，每个epoch结束后，worker进程会被销毁，下一个epoch再重新创建。这有额外开销。设为True后，worker会一直存活，省去了反复创建销毁的时间。我实测过，能节省5%-10%的训练时间。

知识体系总览

下面这张图，把数据加载优化的核心逻辑串起来了：

嗯，这张图其实把整个流程说清楚了。数据从CPU预处理，经过内存缓冲，最后到GPU计算。每个环节都有优化空间，而我们的目标就是让这三个环节流水线起来，不让任何一个环节成为瓶颈。