数据加载与预处理优化：DataLoader调优、数据增强加速、内存映射文件、TFRecord/WebDataset
训练一个模型，GPU算力再强，如果数据喂不上去，那就是白搭。我见过太多团队，花大价钱买了A100，结果GPU利用率只有30%，一查瓶颈全在数据加载上。说白了，数据加载就是整个训练流程的「水管」，水管细了，再大的水泵也白费。
这一章，咱们就聊聊怎么把这根水管加粗。我会从DataLoader调优讲起，再到数据增强的加速技巧，最后聊聊内存映射和TFRecord/WebDataset这些高级玩法。
DataLoader调优：别让CPU闲着
PyTorch的DataLoader，大家都会用。但你真的用对了吗？我刚开始做训练时，就犯过一个低级错误——把num_workers设成了0。结果呢？GPU每算完一个batch，就得干等着CPU去加载下一个，利用率直接掉到20%。
调优DataLoader，核心就几个参数：
| 参数 | 推荐值 | 说明 |
|---|---|---|
| num_workers | 4~8（或CPU核心数） | 子进程数量，不是越大越好 |
| prefetch_factor | 2~4 | 预取多少个batch |
| pin_memory | True | 锁页内存，加速CPU→GPU传输 |
| batch_size | 根据显存调整 | 越大吞吐越高，但注意显存 |
核心原则：让CPU在GPU计算的同时，提前准备好下一个batch的数据。这样GPU永远不用等。
举个例子，我习惯这样配置：
dataloader = DataLoader(
dataset,
batch_size=64,
shuffle=True,
num_workers=8,
pin_memory=True,
prefetch_factor=4,
persistent_workers=True # 避免每个epoch重建worker
)这里有个坑——persistent_workers=True。如果不设这个，每个epoch结束worker会重启，浪费大量时间。我曾经在一个大模型训练中，因为这个参数没设，每个epoch多花了30秒在worker初始化上。30秒看着不多，但训练100个epoch就是50分钟啊。
注意：num_workers不是越大越好。我试过设成32，结果CPU上下文切换开销太大，反而变慢了。一般设成CPU核心数的一半到全部，具体得试。
数据增强加速：别在CPU上做太多
数据增强是提升模型泛化能力的关键，但也是性能杀手。你想想看，每张图片都要做随机裁剪、翻转、颜色抖动……这些操作如果全在CPU上做，GPU就得干等着。
我的经验是：能离线做的，绝不在线做。比如随机裁剪、翻转这些，可以提前生成好，存成增强后的版本。训练时直接加载就行。
但有些增强必须在线做，比如MixUp、CutMix这些跟batch相关的操作。这时候怎么办？
- 用GPU做增强：比如用PyTorch的
torchvision.transforms，或者用NVIDIA的DALI库。DALI可以直接在GPU上做图像解码和增强，速度能快好几倍。 - 用Albumentations：这个库比torchvision的transform快很多，尤其是复杂的增强组合。我做过对比，同样的随机裁剪+颜色抖动，Albumentations能快30%以上。
- 减少不必要的增强：不是所有增强都对模型有帮助。我见过有人给ImageNet训练加了10种增强，结果模型反而过拟合了。嗯，这里要注意，增强不是越多越好。
小技巧：如果必须用CPU做增强，可以考虑把增强操作放到DataLoader的worker里并行做。每个worker独立处理自己的数据，这样能充分利用多核CPU。
内存映射文件：大数据的救星
当你的数据集大到几百GB甚至几TB时，传统的数据加载方式就不行了。全部加载到内存？内存不够。每次从磁盘读？太慢。
这时候，内存映射文件（Memory-Mapped File）就派上用场了。说白了，就是把磁盘上的文件映射到虚拟内存地址空间，让操作系统帮你管理数据的加载和卸载。你访问数据时，就像在访问内存一样，但实际上数据是按需从磁盘加载的。
我曾在一次视频训练任务中用过这个技巧。数据集有500GB，全是高清视频帧。如果用传统方式加载，光读取索引就要花10分钟。改用内存映射后，启动时间降到了10秒以内，而且训练过程中几乎没有IO等待。
PyTorch里用起来很简单：
import numpy as np
import torch
# 创建内存映射文件
data = np.memmap('data.bin', dtype='float32', mode='r', shape=(100000, 3, 224, 224))
# 训练时直接索引
for idx in range(100000):
sample = torch.from_numpy(data[idx])
# 处理sample...优点：支持超大数据集、启动快、内存占用低。
缺点：随机访问时如果数据不连续，会有大量缺页中断，反而变慢。
所以，内存映射文件最适合顺序访问的场景。如果你的训练是随机打乱的，那就要小心了——每次随机访问都可能触发磁盘IO。我建议的做法是：先用内存映射把数据顺序读一遍，缓存到内存里，然后再做随机打乱。
TFRecord与WebDataset：结构化数据格式
说到大数据加载，就不得不提TFRecord和WebDataset。这两个都是专门为高性能训练设计的数据格式。
TFRecord是TensorFlow的亲儿子，但PyTorch也能用。它的核心思想是把多个样本打包成一个文件，减少小文件IO的开销。你想想看，如果每个样本都是一个单独的图片文件，那读取100万张图片就要打开100万次文件，光文件系统的开销就够受的。
TFRecord的用法：
import tensorflow as tf
# 写入TFRecord
with tf.io.TFRecordWriter('data.tfrecord') as writer:
for image, label in dataset:
example = tf.train.Example(features=tf.train.Features(feature={
'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image.tobytes()])),
'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))
}))
writer.write(example.SerializeToString())
# 读取TFRecord
raw_dataset = tf.data.TFRecordDataset('data.tfrecord')
for raw_record in raw_dataset:
example = tf.train.Example()
example.ParseFromString(raw_record.numpy())
# 解析数据...WebDataset是更新的方案，它基于tar格式，但做了优化。每个tar文件里包含多个样本，每个样本可以是任意格式（图片、文本、音频等）。WebDataset最大的优势是流式读取——你不需要把所有数据都加载到内存，而是边读边训练。
我个人更推荐WebDataset，原因有三：
- 兼容性好：任何语言、任何框架都能读tar文件
- 支持分布式：天然适合多节点训练，每个节点读不同的tar文件
- 速度快：顺序读取tar文件，磁盘IO效率极高
避坑指南：我曾经在WebDataset里犯过一个错——把tar文件压缩了。结果每次读取都要解压，反而比不压缩还慢。记住，WebDataset追求的是读取速度，不是存储空间。除非你的网络带宽是瓶颈，否则别压缩。
知识体系总览
下面这张图，是我对本章知识体系的总结。你可以把它当作一个决策树，遇到数据加载问题时，按图索骥就行。
总结一下，数据加载优化的核心思路就四个字：减少等待。让CPU、GPU、磁盘三者各司其职，谁也不等谁。具体怎么做？
- DataLoader调优：num_workers、pin_memory、prefetch_factor，这三个参数调好了，CPU利用率能翻倍。
- 数据增强加速：能离线做的离线做，不能离线做的用GPU做，或者用更快的库。
- 内存映射文件：适合超大数据集，顺序读取时效率极高。
- TFRecord/WebDataset：结构化数据格式，减少小文件IO，支持流式读取。
嗯，这些技巧我都在实际项目中验证过。下次你遇到训练速度上不去，别急着换显卡，先看看数据加载这块有没有优化空间。很多时候，调好DataLoader比换一张A100还管用。