4、数据加载优化：DataLoader多进程设置、prefetch_factor调优、内存映射文件、数据预处理流水线、避免IO瓶颈
数据加载，说白了就是喂饭给GPU吃。
你想想看，GPU算得再快，如果数据送不上去，它就得干等着。我见过太多团队，花大价钱买了8卡、16卡机器，结果训练吞吐上不去，一查原因——IO瓶颈。GPU利用率才30%，剩下的时间都在等数据。
这一章，我们就来聊聊怎么把数据加载这个环节彻底打通。
4.1 DataLoader多进程设置：别让主进程干苦力
PyTorch的DataLoader默认是单进程加载数据。这意味着什么？意味着每次要取一个batch，主进程都得亲自去磁盘读数据、做预处理。GPU在等你，CPU也在忙，效率极低。
我个人习惯，只要数据量超过1万条，就一定会开多进程。核心参数就一个：num_workers。
核心原则：num_workers通常设置为CPU核心数的1/2到2/3。不是越多越好，太多反而会带来进程切换开销。
from torch.utils.data import DataLoader
# 我常用的配置
dataloader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=8, # 根据CPU核心数调整
pin_memory=True, # 加速CPU到GPU的数据传输
prefetch_factor=2 # 预取因子，后面会讲
)我在项目中遇到过一个问题：开了16个worker，结果训练反而变慢了。为什么？因为磁盘IO成了瓶颈，16个进程都在抢着读同一个磁盘，互相打架。后来我降到8个worker，速度反而上来了。
注意：如果使用多进程，记得把数据加载的逻辑放在 if __name__ == '__main__': 下面，否则Windows系统会报错。
4.2 prefetch_factor调优：让数据提前就位
prefetch_factor，这个参数很多人会忽略。它控制每个worker预取多少个batch的数据。
默认值是2，意味着每个worker会提前准备2个batch。如果你的数据预处理很耗时，比如要做大量的图像增强、文本tokenization，那2可能不够。
# 对于预处理重的任务，我建议调大prefetch_factor
dataloader = DataLoader(
dataset,
batch_size=64,
num_workers=8,
prefetch_factor=4, # 每个worker预取4个batch
pin_memory=True
)我曾经做过一个实验：同样的模型，同样的数据，prefetch_factor从2调到4，训练吞吐提升了12%。原因很简单——GPU在算当前batch的时候，worker们已经在准备后面4个batch了，流水线完全跑满。
调优技巧：prefetch_factor * num_workers 不要超过总内存的30%。否则数据预取太多，内存爆了，反而触发swap，速度暴跌。
4.3 内存映射文件：省掉一次拷贝
内存映射文件，英文叫mmap。说白了，就是把磁盘文件直接映射到进程的虚拟地址空间。读文件的时候，不需要先读到用户态缓冲区，再拷贝到进程内存——直接访问映射区域就行。
对于大文件，比如几百GB的文本数据、图像数据集，mmap能省掉一次内存拷贝，IO效率提升明显。
import numpy as np
import mmap
# 创建内存映射文件
with open('large_dataset.bin', 'r+b') as f:
# 映射整个文件
mmapped_file = mmap.mmap(f.fileno(), 0)
# 直接当作numpy数组使用
data = np.frombuffer(mmapped_file, dtype=np.float32).reshape(-1, 784)
# 用完记得关闭
mmapped_file.close()嗯，这里要注意：mmap适合只读场景。如果你要频繁修改数据，mmap反而会因为写回磁盘而变慢。
我的经验：在训练GPT类大模型时，我把训练数据全部转成mmap格式，加载速度从原来的每epoch 15分钟降到了3分钟。省下来的时间，够我多跑好几个实验了。
4.4 数据预处理流水线：把计算藏起来
数据预处理，比如图像resize、归一化、文本分词，这些操作如果放在训练循环里做，GPU就得等你。正确的做法是：把这些操作放到DataLoader的worker进程里，和训练并行执行。
这就是数据预处理流水线的核心思想——用CPU的时间换GPU的等待。
class MyDataset(Dataset):
def __init__(self, data_path):
self.data = load_data(data_path)
# 预处理放在这里？不，放在__getitem__里
def __getitem__(self, idx):
sample = self.data[idx]
# 在这里做预处理，worker进程会并行执行
sample = self.preprocess(sample) # resize, normalize, etc.
return sample
def preprocess(self, sample):
# 耗时操作，比如图像增强
sample = 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])
])(sample)
return sample你可能会问：那预处理的结果要不要缓存？我的建议是：如果数据量不大（比如几万张图），可以缓存到内存里。但如果数据量上百万，缓存反而会撑爆内存，得不偿失。
进阶技巧：对于重复使用的预处理结果，可以用LMDB或HDF5格式持久化。训练时直接读预处理好的数据，连预处理这一步都省了。
4.5 避免IO瓶颈：实战中的那些坑
IO瓶颈，说白了就是数据从磁盘到内存、从内存到显存这条路上，哪个环节慢了，整个训练就得等。
我总结了几种常见的IO瓶颈和解决方案：
| 瓶颈类型 | 表现 | 解决方案 |
|---|---|---|
| 磁盘读取慢 | GPU利用率低，CPU等待IO | 换SSD，或者用内存文件系统（tmpfs） |
| CPU预处理慢 | worker进程CPU跑满 | 增加num_workers，优化预处理逻辑 |
| CPU到GPU传输慢 | pin_memory未开启 | 设置pin_memory=True，使用非阻塞传输 |
| 数据格式不合理 | 频繁小文件读取 | 合并成大文件，使用TFRecord或MDS格式 |
我曾经踩过一个坑：训练一个图像分类模型，数据集是10万张JPEG图片，每张几百KB。每次读取都要解压JPEG，CPU全耗在解压上了。后来我把所有图片转成LMDB格式，读取速度提升了5倍。
避坑指南：千万不要在训练循环里做文件IO操作。比如每个step都打开一个文件、读取、关闭——这会让你的训练速度直接腰斩。所有IO操作都应该在DataLoader的worker里完成。
4.6 知识体系总览
下面这张图，是我对数据加载优化整个知识体系的总结。你可以把它当作一个检查清单，看看自己的训练流程里，哪个环节还没优化到位。
数据加载优化，说难不难，说简单也不简单。核心就一句话：让数据等GPU，别让GPU等数据。你只要把上面这几个点都检查一遍，训练吞吐至少能提升30%以上。
好了，这一章就聊到这里。下一章我们继续深入，聊聊分布式训练中的通信优化。