3、数据加载优化（上）：DataLoader 的工作原理，num_workers 的设置技巧，pin_memory 的作用

聊到训练吞吐，很多人第一反应就是调大 batch size、换更快的 GPU。但我得说句实话——数据加载往往是那个被忽视的瓶颈。我见过太多项目，GPU 利用率只有 30%，一看监控，数据加载线程在那边吭哧吭哧地跑，GPU 却在干等着。说白了，你的模型再快，数据喂不进去也是白搭。

DataLoader 到底在干什么？

PyTorch 的 DataLoader，本质上就是个数据搬运工。它负责从硬盘上把数据读出来，做预处理（比如裁剪、归一化），然后打包成 batch 送给 GPU。

但这里有个关键问题：数据加载是 I/O 密集型操作，而模型训练是计算密集型操作。这两者天生就不匹配。你想想看，GPU 一秒能算几百张图，但硬盘读一张图可能就要几十毫秒。如果让 GPU 等着数据加载完再算，那效率得多低？

所以 DataLoader 的设计思路很简单：用多线程/多进程提前把数据准备好。等 GPU 算完当前 batch，下一个 batch 已经躺在内存里了，直接拿过来用就行。

num_workers：到底设多少合适？

这是最常被问的问题，也是我踩坑最多的地方。num_workers 控制的是数据加载的子进程数量。设少了，数据供不上；设多了，系统资源被抢光，反而更慢。

我个人习惯的调优思路是这样的：

1. 先试 0：也就是在主进程里加载数据。这能帮你确认数据加载本身是不是瓶颈。如果 GPU 利用率已经很高了，那就不用折腾了。
2. 从 2 开始往上加：一般建议从 num_workers = 2 开始，观察 GPU 利用率的变化。每次翻倍，直到利用率不再提升。
3. 注意上限：我一般不会超过 CPU 核心数。比如你机器有 8 核，设到 6 或 7 就差不多了。留一个核心给系统和其他进程。

我曾经遇到过一个问题：把 num_workers 设得太大，结果每个 worker 都在抢内存，系统开始频繁 swap，训练速度反而下降了 30%。后来我限制到 4 个 worker，速度立马就上来了。嗯，这里要注意：不是越多越好。

pin_memory：这个参数到底有什么用？

说实话，我刚开始用 PyTorch 时完全忽略了 pin_memory。直到有一次，一个同事问我：「你设了 pin_memory=True 吗？」我才去研究了一下。

它的作用其实很简单：把数据固定在内存的某个特殊区域里，让 GPU 可以直接通过 DMA 访问。如果没有 pin_memory，数据从 CPU 传到 GPU 需要先拷贝到页锁定内存，多了一步操作。

说白了，pin_memory=True 就是给数据传输开了个「绿色通道」。数据从 CPU 到 GPU 的速度会快很多。

但要注意一点：pin_memory 会占用更多的物理内存。因为页锁定内存不能被换出到磁盘。如果你的内存比较紧张（比如只有 8GB），可能需要权衡一下。

一个完整的配置示例

from torch.utils.data import DataLoader

# 我个人常用的配置
dataloader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,        # 根据 CPU 核心数调整
    pin_memory=True,      # 内存够就开着
    prefetch_factor=2,    # 每个 worker 预取 2 个 batch
    persistent_workers=True  # 复用 worker 进程
)

这里我额外加了两个参数：prefetch_factor 和 persistent_workers。前者控制每个 worker 预取多少个 batch，后者让 worker 在 epoch 之间不销毁重建。这两个配合起来，能进一步减少数据加载的延迟。

知识体系总览

下面这张图总结了 DataLoader 的核心逻辑和优化方向：

好了，这一章的内容就到这里。数据加载优化是个系统工程，今天讲的是 DataLoader 的基础配置。下一章我们会深入聊聊更高级的技巧，比如内存映射、数据预处理流水线等。但先把今天这些基础打好，你的训练吞吐就能提升一大截。