4. 数据加载优化：DataLoader 工作原理、num_workers 调优、pin_memory 的作用、prefetch_factor 设置
聊到显存优化，很多人第一反应是改模型结构、调 batch size。但说实话，我见过太多人把精力全花在模型上，结果数据加载成了瓶颈——GPU 在那儿干等着，利用率才 30%。
数据加载优化，说白了就是让 GPU 别闲着。你想想看，模型算得再快，数据送不上去也是白搭。这一章，我就带你彻底搞懂 PyTorch 的 DataLoader，以及它背后那些能帮你省显存、提速度的关键参数。
核心观点：数据加载优化的本质，是用 CPU 的预处理能力，去掩盖 IO 延迟，从而让 GPU 始终处于满负荷计算状态。
4.1 DataLoader 的工作原理
DataLoader 不是简单地把数据一批批扔给模型。它内部有一套完整的流水线机制。我个人习惯把它拆成三个环节：
- 采样器（Sampler）：决定从数据集中取哪些样本。比如随机采样、顺序采样、分布式采样。
- 读取与预处理：把样本从硬盘读到内存，然后做 transform（归一化、裁剪、增强等）。
- 组装成 batch：把多个样本拼成一个张量，送到 GPU。
默认情况下，这些步骤都在主进程里串行执行。但问题来了——如果每张图片都要从硬盘读、做数据增强，CPU 处理时间可能比 GPU 计算时间还长。这就是为什么我们需要多进程。
小提示：DataLoader 的 collate_fn 参数可以自定义 batch 组装逻辑。比如处理变长序列时，你可以在这里做 padding，而不是在 transform 里做。
4.2 num_workers 调优
num_workers 是 DataLoader 里最直接的加速参数。它决定了用几个子进程来并行加载数据。
我曾经在一个图像分类项目里，把 num_workers 从 0 调到 4，训练速度直接翻了一倍。但注意，不是越大越好。
| num_workers 值 | 适用场景 | 潜在问题 |
|---|---|---|
| 0 | 调试阶段、数据量极小 | GPU 利用率低 |
| 1~4 | 单机单卡、CPU 核心数较少 | 可能仍有等待 |
| 4~8 | 单机多卡、CPU 资源充足 | 内存占用上升 |
| 超过 CPU 核心数 | 不推荐 | 进程切换开销 > 收益 |
调优时，我一般遵循一个经验法则：num_workers 设为 CPU 核心数的一半到三分之二。比如 8 核 CPU，设 4 或 6 比较稳妥。
避坑指南：我曾经在 Windows 上把 num_workers 设得过高，结果程序直接卡死。原因是 Windows 的进程创建机制和 Linux 不同，子进程太多会导致内存爆炸。如果你用 Windows，建议 num_workers 不超过 4。
4.3 pin_memory 的作用
pin_memory 这个参数，很多人直接设成 True 就完事了。但你真的理解它在干什么吗？
说白了，pin_memory 是把数据固定在 CPU 的物理内存中，不让它被换到虚拟内存（swap）里。这样做的好处是：当 GPU 要从 CPU 内存拷贝数据时，可以直接通过 DMA（直接内存访问）传输，速度更快。
没有 pin_memory 时，数据要先从可分页内存拷贝到固定内存，再传给 GPU。多了一步拷贝，延迟就上去了。
我的建议：只要你的显存不是特别紧张，pin_memory=True 几乎总是有益的。它不会增加太多内存开销，但能显著提升数据传输速度。
但注意一点：pin_memory 只在你用 .to(device) 或 .cuda() 把数据搬到 GPU 时才有用。如果你全程在 CPU 上训练，设不设都一样。
4.4 prefetch_factor 设置
prefetch_factor 是 PyTorch 1.7 之后引入的参数，很多人不知道它的存在。它控制每个 worker 预先加载多少个 batch 的数据。
默认值是 2，意思是每个 worker 会提前加载 2 个 batch 的数据到内存里。这样当 GPU 处理当前 batch 时，下一个 batch 已经在内存里等着了。
我遇到过这样一个场景：数据增强特别复杂，每张图要随机裁剪、旋转、颜色抖动，CPU 处理时间很长。这时候把 prefetch_factor 从 2 调到 4，GPU 的等待时间明显减少了。
| prefetch_factor | 效果 | 风险 |
|---|---|---|
| 1 | 最小内存占用 | GPU 可能等待 |
| 2（默认） | 平衡内存与速度 | 适合大多数场景 |
| 4 或更高 | 减少 GPU 空闲 | 内存占用翻倍 |
调优技巧：你可以用 nvidia-smi 观察 GPU 利用率。如果利用率经常低于 80%，说明数据加载是瓶颈。这时候优先调大 num_workers，再考虑 prefetch_factor。
4.5 知识体系总览
下面这张图，我把 DataLoader 优化的核心逻辑画出来了。你可以看到，数据从硬盘到 GPU 的每一步，都有对应的优化手段。
4.6 实战代码示例
下面是一个完整的 DataLoader 配置示例。你可以直接拿去用，改改路径就行。
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
# 自定义数据集
class MyDataset(Dataset):
def __init__(self, data_list, transform=None):
self.data = data_list
self.transform = transform
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
sample = self.data[idx]
if self.transform:
sample = self.transform(sample)
return sample
# 数据增强
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# 创建 DataLoader —— 优化配置
dataloader = DataLoader(
dataset=MyDataset(data_list, transform=transform),
batch_size=64,
shuffle=True,
num_workers=4, # 根据 CPU 核心数调整
pin_memory=True, # 加速 CPU→GPU 传输
prefetch_factor=2, # 每个 worker 预取 2 个 batch
persistent_workers=True # 避免每个 epoch 重建 worker
)注意：persistent_workers=True 这个参数容易被忽略。它让 worker 进程在 epoch 之间保持存活，而不是每次重新创建。如果你的数据集不大，这个参数能省掉不少进程启动开销。
4.7 调优实战经验
最后，分享几个我踩过的坑：
- num_workers 设太大，内存爆了——每个 worker 都会复制一份数据预处理逻辑，内存占用是线性增长的。我建议你边调边看
htop，别让内存超过 80%。 - pin_memory 在 CPU 训练时没用——如果你只在 CPU 上跑，设了也白设。但如果你不确定，设成 True 也没坏处。
- prefetch_factor 不是越大越好——它本质是用内存换时间。如果你的内存本来就紧张，设成 1 或 2 就够了。
嗯，数据加载优化这块，说白了就是让 CPU 和 GPU 各司其职，别让任何一方闲着。你按我上面说的顺序调，基本不会出大问题。