数据加载优化：DataLoader工作原理、num_workers设置、prefetch_factor调整、内存映射文件
训练跑得慢，很多时候不是GPU在偷懒，而是CPU没把饭喂到嘴边。
我见过太多团队，花大价钱买了A100，结果GPU利用率只有30%。一看监控，数据加载那块儿卡得死死的。说白了，数据加载就是训练管道的入口，入口堵了，后面再快也没用。
DataLoader工作原理：它到底在干嘛？
PyTorch的DataLoader，本质上就是个数据搬运工。它负责从硬盘上把数据读出来，做预处理，然后打包成batch送给模型。
它的核心逻辑其实很简单：
- Sampler：决定取哪些样本，比如随机取还是顺序取
- Dataset：给定一个索引，返回对应的样本数据
- Collate_fn：把多个样本拼成一个batch
嗯，这里要注意，很多人以为DataLoader是并行工作的。其实不是。单进程模式下，它就是一个接一个地读数据，CPU在等硬盘，GPU在等CPU，整个链条都在等。
核心要点：DataLoader默认是单进程模式。想要快，必须开多进程。
num_workers设置：开几个搬运工最合适？
num_workers就是开几个子进程来帮你读数据。我个人的习惯是，先试CPU核心数，再往下调。
举个例子，一台机器有16个物理核，32个逻辑核。我会先设num_workers=8，然后看GPU利用率。如果GPU还在等，就加到12。如果CPU被打满了，就降到6。
为什么不是越多越好？
我曾经踩过一个坑。有一次我把num_workers设成了32，想着反正核多。结果训练反而变慢了。为什么呢？因为进程太多，操作系统在进程切换上花的时间比干活还多。而且每个进程都要复制数据，内存带宽成了瓶颈。
我的经验公式：
- CPU密集型预处理（如图像增强）：num_workers = CPU核心数 / 2
- IO密集型（简单读取）：num_workers = CPU核心数
- 千万别超过CPU核心数的2倍，否则适得其反
prefetch_factor调整：提前多备几份饭
prefetch_factor这个参数，很多人不知道。它控制每个worker提前加载多少个batch。
默认值是2，意思是每个worker会提前加载2个batch的数据在内存里等着。GPU吃完一个，马上就能拿下一个。
你想想看，如果prefetch_factor=1，worker刚把数据喂给GPU，就得立刻去读下一批。万一硬盘慢了，GPU就得空转。prefetch_factor大一点，相当于给GPU建了个缓冲区。
我一般怎么调？
- 数据在SSD上：prefetch_factor=2或3就够了
- 数据在机械硬盘或网络存储上：我会调到4或5，给足缓冲
- 内存够大但CPU吃紧：调大prefetch_factor，让worker少干活
注意：prefetch_factor不是越大越好。调太大，内存占用会飙升。我曾经在一个96GB内存的机器上，把prefetch_factor调到8，结果内存直接爆了。嗯，那次排查了好久才发现是它搞的鬼。
内存映射文件：让硬盘假装是内存
内存映射文件，英文叫memory-mapped file，简称mmap。这玩意儿说白了，就是把硬盘上的文件映射到虚拟内存地址空间里。程序读写这个文件，就像读写内存一样快。
为什么它能加速数据加载？
- 省去了read()系统调用的开销
- 操作系统会自动做缓存，热数据留在内存里
- 多个进程可以共享同一块映射区域，省内存
我在项目中遇到过一个大模型训练任务，数据集有2TB，全是高清图片。用普通的ImageFolder加载，一个epoch要跑6个小时。后来改成mmap模式，直接把整个数据集映射到内存地址空间，加载时间降到了40分钟。
怎么用？PyTorch的Dataset里，用numpy的memmap或者直接操作文件描述符就行。代码大概长这样：
import numpy as np
import torch
from torch.utils.data import Dataset
class MmapDataset(Dataset):
def __init__(self, data_path, label_path, shape, dtype=np.float32):
# 内存映射文件，不实际加载到内存
self.data = np.memmap(data_path, dtype=dtype, mode='r', shape=shape)
self.labels = np.memmap(label_path, dtype=np.int64, mode='r')
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
# 这里读取时，操作系统会自动做页缓存
x = torch.from_numpy(self.data[idx].copy())
y = self.labels[idx]
return x, y小技巧：如果数据集太大，内存装不下，可以只映射一部分。或者用多个mmap文件，按需加载。我习惯把数据按epoch分片，每个片一个mmap文件，这样既快又省内存。
知识体系总览
下面这张图，把数据加载优化的几个关键点串起来了。你可以看到，从硬盘到GPU，每一步都有优化的空间。
实战建议：怎么搭一套高效的数据管道？
说了这么多，给你一套我常用的配置模板：
from torch.utils.data import DataLoader
dataloader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=8, # 根据CPU核心数调整
prefetch_factor=3, # 提前加载3个batch
pin_memory=True, # 锁页内存，加速CPU到GPU传输
persistent_workers=True # 保持worker存活，减少创建销毁开销
)额外提醒：pin_memory=True一定要开。它把数据固定在内存里，GPU可以直接通过DMA访问，不用CPU中转。这个开关能再给你10%~20%的提速。
最后说一句，数据加载优化没有银弹。你得根据你的硬件配置、数据集大小、预处理复杂度来调。我建议你先用默认参数跑一遍，然后逐个调优，每次只改一个参数，看GPU利用率的变化。
嗯，今天就聊到这儿。数据加载这块儿，说白了就是让CPU和GPU之间的管道更通畅。管道通了，训练自然就快了。