数据加载管线：从磁盘到GPU的最后一公里
数据加载，说白了就是训练流程里的「送餐员」。GPU再快，如果数据送不到嘴边，也只能干等着。我见过太多团队，模型架构调得飞起，结果训练速度被DataLoader卡死——GPU利用率只有30%，显存倒是没爆，但时间全浪费在等数据上了。
这一章，咱们就聊聊数据加载管线的核心：DataLoader怎么工作的，prefetch_factor和num_workers怎么调，以及mmap这个「黑科技」到底能省多少事。
DataLoader的工作原理：别把它当黑盒
PyTorch的DataLoader，本质上是一个多进程数据生产流水线。它的工作流程可以拆成三步：
- 采样器（Sampler）：决定从数据集的哪个位置取数据
- 读取器（Reader）：从磁盘加载原始数据（图片、文本等）
- 处理器（Processor）：做数据增强、归一化、转Tensor等操作
这三步默认是串行的。但如果你设置了num_workers > 0，读取和处理器就会跑到子进程里，主进程只管拿结果。
这里有个关键点：每个worker进程都会复制一份数据集对象。如果你的数据集里存了很大的缓存（比如把所有图片路径都加载到内存），那每个worker都会有一份副本——显存就是这么悄悄被吃掉的。
核心结论：DataLoader不是「多线程」，而是「多进程」。每个worker是独立进程，有自己的内存空间。这意味着：
- 数据加载可以并行
- 但内存开销也会翻倍
- worker之间不共享任何状态
prefetch_factor调优：预取多少才够？
prefetch_factor控制的是每个worker预取多少个batch的数据。默认值是2，意思是每个worker会提前加载2个batch放到队列里。
我个人习惯把这个参数和num_workers一起调。为什么？因为预取队列的总大小 = num_workers × prefetch_factor。如果这个值太大，内存会爆；太小，GPU又可能饿死。
| 场景 | num_workers | prefetch_factor | 效果 |
|---|---|---|---|
| 小数据集（内存能放下） | 4-8 | 2-4 | 预取足够，GPU不饿 |
| 大数据集（磁盘IO瓶颈） | 8-16 | 4-8 | 掩盖IO延迟 |
| 显存紧张 | 2-4 | 1-2 | 减少内存占用 |
我在项目中遇到过一个问题：用SSD加载ImageNet，num_workers设了16，prefetch_factor设了8。结果训练到一半，内存直接飙到80GB——每个worker预取了8个batch，每个batch 256张图，每张图解码后几十MB...你算算这得多少内存？
避坑指南：prefetch_factor不是越大越好。它和batch_size、图片尺寸、解码格式强相关。我曾经把prefetch_factor从4调到8，结果训练速度没提升，内存先爆了。
num_workers最佳实践：不是越多越好
很多人觉得num_workers设得越大，数据加载越快。其实不然。worker太多会带来两个问题：
- 进程切换开销：CPU在多个worker之间频繁切换，反而降低了效率
- 内存爆炸：每个worker都有自己的数据副本
我的经验公式是这样的：
- CPU密集型任务（如图像解码+数据增强）：num_workers = CPU核心数 / 2
- IO密集型任务（如从网络加载数据）：num_workers = CPU核心数
- 混合型：num_workers = CPU核心数 × 0.75
举个例子。我调过一个OCR模型，数据加载主要是读图片+做旋转扭曲。CPU是32核，我一开始设了32个worker，结果CPU利用率只有60%，内存却用了40GB。后来降到16个worker，CPU利用率反而到了85%，内存降到20GB，训练速度还快了10%。
小技巧：用nvidia-smi和htop同时观察GPU利用率和CPU利用率。如果GPU利用率低于80%但CPU已经跑满，说明worker不够；如果GPU利用率低但CPU还有余量，说明worker太多或者IO是瓶颈。
内存映射文件（mmap）：省内存的终极方案
mmap，全称memory-mapped file。说白了就是把磁盘文件直接映射到进程的虚拟地址空间。你读写这个文件，就像读写内存一样——操作系统会在背后帮你做磁盘和内存之间的数据交换。
为什么这对训练有好处？
- 省内存：不需要一次性把整个数据集加载到内存
- 省时间：避免了从磁盘到用户空间的拷贝
- 共享内存：多个worker可以共享同一个mmap映射，不会重复占用内存
PyTorch的torch.utils.data.Dataset里，你可以这样用mmap：
import mmap
import numpy as np
class MMapDataset(Dataset):
def __init__(self, file_path, shape, dtype=np.float32):
self.file = open(file_path, 'rb')
self.mmap = mmap.mmap(
self.file.fileno(),
0, # 映射整个文件
access=mmap.ACCESS_READ
)
self.shape = shape
self.dtype = dtype
self.data_size = np.prod(shape) * np.dtype(dtype).itemsize
def __len__(self):
return len(self.mmap) // self.data_size
def __getitem__(self, idx):
start = idx * self.data_size
end = start + self.data_size
# 直接从mmap读取，不拷贝到Python内存
data = np.frombuffer(
self.mmap[start:end],
dtype=self.dtype
).reshape(self.shape)
return torch.from_numpy(data)
这里有个关键点：np.frombuffer返回的是视图，不是拷贝。也就是说，你读取数据时，并没有把数据从mmap复制到Python堆里——这能省下大量内存。
实战效果：我在一个100GB的文本数据集上做过对比：
- 普通加载：内存占用120GB（数据集+缓存+worker副本）
- mmap加载：内存占用8GB（只有索引和元数据）
- 训练速度：mmap版本比普通版本快15%
数据加载管线的完整流程图
下面这张图展示了数据从磁盘到GPU的完整路径，以及各个参数的作用位置：
从这张图你能看到：mmap让多个Worker共享同一份磁盘映射，避免了重复加载；预取队列则像一个缓冲池，让GPU随时有数据可用。这两个机制配合好了，数据加载基本不会成为瓶颈。
最后说一句：调优数据加载管线，核心思路就四个字——「减少等待」。要么让数据提前准备好（预取），要么让数据加载更快（mmap），要么让加载过程不阻塞（多进程）。具体用哪个，取决于你的瓶颈在哪。
想要更多资料，联系可可