数据加载成为训练瓶颈时，核心解法就是让数据管线的每一层都跑起来，用多进程并行和预取把等待时间抹平。模型训练时GPU空转、显存吃不满，根源不在于算力不够，而是数据供给跟不上，这个问题在CV、NLP、推荐系统的大规模训练里相当常见,搞明白优化方向比盲目加显卡更管用。
先确认瓶颈到底在不在数据加载这段
不少同学一看GPU利用率低就怀疑数据加载，其实先要验证这个假设。用PyTorch Profiler或者TensorBoard看一眼训练时间线，如果DataLoader的耗时占比明显高于前向和反向，或者GPU的Utilization长期偏低而CPU跑满,那基本就能锁定问题。
一个快速验证办法：把DataLoader换成IterableDataset临时跑一遍，或者干脆用torch.utils.data.DataLoader的num_workers=0跑一次对比，如果速度反而上去了，说明多进程开销在捣乱；如果还是很慢,就得去查存储设备本身的读取能力。
行业共识认为，数据管线的理想状态是GPU算完一个batch的时间，刚好等于数据管线准备好下一个batch的时间。 任何时候出现一方等待另一方,都是需要优化的信号。
从PyTorch DataLoader参数入手，成本最低
num_workers设置多少合适
这个参数经常被拍脑袋设置，它决定了有几个子进程同时帮你做数据预处理和加载。设置太小，主进程忙不过来，设置太大，进程间通信和内存复制反而拖慢速度。
- 单机场景下，从
4开始试,逐步加到CPU物理核心数的一半左右。 - GPU训练时，每个GPU配
2~4个worker通常是够用的。 - 如果发现内存占用过高或者频繁报
File descriptor错误,那就是开多了。
prefetch_factor：预取几批数据
prefetch_factor决定了每worker在内存里预存几个batch的数据。 默认值是2，意味着当前batch喂给模型的时候，后面已经有2个batch在排队，把这个值调大到4或者8，能有效遮挡掉数据读取的尖峰延迟,尤其是在文件系统抖动比较明显的场景下。
pin_memory和non_blocking：绕过CPU到GPU的搬运
把pin_memory=True打开，配合.cuda(non_blocking=True)，数据从CPU内存搬到GPU显存的路径会变成异步的，不占用主线程的计算时间，这个操作几乎是无本万利，对于每次batch里图片或文本规模较大的情况，能缩短不少搬运耗时。
关掉多余的数据增强
训练时在CPU上做随机裁剪、色彩抖动这些操作，讲道理是必要的，但如果增强逻辑写得太重，比如每张图片都要做高斯模糊加随机擦除，那CPU就成了新的堵点，常见的做法是把一部分增强操作搬到GPU上做，用torchvision.transforms里的GPU版本，或者干脆在TensorFlow里用tf.image配合prefetch到GPU设备上执行。
存储设备和文件格式，通常是被忽略的大头
机械硬盘到NVMe SSD的差距有多大
数据加载从HDD搬到NVMe SSD，所带来的训练提速往往比换一张显卡还要明显。 有些大厂内部做过统计，在CV训练任务里，仅仅是存储设备的升级，就能把数据读取阶段延迟缩短一个数量级，如果预算有限，优先把训练数据放到NVMe盘上,而不是先加内存。
TFRecord和LMDB这类格式能救场
小文件散落成几万个图片路径，每次都要open/close文件描述符，慢是必然的。把数据打包成TFRecord（TensorFlow的二进制格式）或者LMDB（轻量级数据库）格式，顺序读取大文件,能明显减少IO次数。
| 存储方案 | 适合场景 | 主要优势 | 劣势 |
|---|---|---|---|
| 原始图片+路径读取 | 数据量小、实验期 | 简单直观 | 小文件IO极慢 |
| TFRecord | TensorFlow训练 | 顺序读取，天然适合数据管道 | 打包麻烦，改格式要时间 |
| LMDB | PyTorch训练 | 读取快，支持随机访问 | 需要额外维护缓存 |
| WebDataset | 超大规模分布式训练 | 流式读取，配合多节点友好 | 生态相对新 |
内存映射（mmap）是省钱又高效的路子
把数据文件映射到虚拟内存，省掉一次用户态的read调用。 PyTorch的torch.load在加载整个数据集时其实就能用mmap模式，业界不少开源项目在热启动时也会用这个思路，对单机训练来说,这是性价比很高的优化。
数据预处理能不能提前做好，别在训练时现算
特征缓存与向量化：弄一次，存起来
很多情况下，数据加载慢不是因为读文件慢，而是在预处理阶段反复做同样的事情，比如NLP任务的分词、ID映射，每次epoch都重新做一遍，纯属浪费，把这些预处理结果缓存到.npy或者parquet格式里，训练时只读不进CPU,效果立竿见影。
- 分词和词表映射提前做好，保存成token ID数组。
- 图片的resize和归一化在离线阶段完成,训练时只做张量化。
- 特征工程结果可以用
joblib或pickle落地,下次直接加载。
训练数据加载慢怎么办：试试缓存到内存文件系统
如果内存足够大，把训练数据整个丢进/dev/shm（共享内存文件系统）或者用memfd_create，读取速度跟内存带宽一个量级，比任何SSD都快得多。 注意别把系统内存撑爆，配合num_workers和prefetch调优,这套组合拳能解决绝大多数单机训练场景的瓶颈。
分布式训练数据加载的额外功课
每个节点只读自己那份数据
用TensorFlow的tf.data配合distribute策略，或者PyTorch的DistributedSampler，都能做到数据分片。最怕的是所有节点从头到尾读同一份完整数据，还嫌弃网络带宽不够用。 行业共识是让每个节点访问本地存储的数据分片,避免跨节点拷贝。
用GPU直接做数据解码
NVIDIA的DALI库能把JPEG解码和图像增强放到GPU上跑，CPU彻底退出数据管线的主力位置。 配合--prefetch队列深度调高，吞吐量能上来一大截，这项技术在推荐系统和搜索场景的CTR预估模型训练里用得非常多,因为稀疏特征的数据预处理太重了。
避免频繁创建和销毁DataLoader对象
每个epoch结束都重新DataLoader(...)，意味着子进程重建、内存重新分配，开销不小。把DataLoader定义在循环外，用iterator重制的方式复用。
for epoch in range(epochs):
for batch in train_loader:
...这样train_loader的子进程池保持存活,省掉反复启动的时间。
实测里的几个判断标准
判断优化有没有到位,看两个指标就行：
- Time to train（TTT）缩短了多少。 训练不同规模的数据集,记录完整跑完的墙钟时间。
- GPU利用率稳定在什么水平。 如果稳定图片高位,说明数据管线已经不再是短板。
如果调完DataLoader还是慢，打开top看CPU的sys占用是不是特别高，如果是，那多半是内核态IO太频繁,按上面的存储优化方向去搞。
Q&A：数据加载成为训练瓶颈时的优化方向
除了换硬件，还有什么低成本优化手段？
优先做软件层面的调整：把num_workers从默认值往上调、打开prefetch和pin_memory、数据格式从散图打包成TFRecord或LMDB、预处理结果离线缓存，据PyTorch官方文档说明,这几项组合起来通常能让数据管线运行时间缩短数倍。
分布式训练里，数据加载和单机有什么不同？
分布式场景下，每个节点只能看到自己负责的那份数据分片，额外多了一个分片策略和网络传输的问题，用ShardedDataset或者WebDataset按tar包切分，节点内直接读本地文件，不再走共享存储的网络路径,这是比加大带宽更实在的解法。
数据增强太耗时，有没有办法既保留效果又不拖慢训练？
把部分增强操作放到GPU上做，或者使用NVIDIA DALI的GPU算子，把解码、resize、随机裁剪全部卸载到GPU端，这样一来CPU只做轻量级的调度工作，GPU也基本有空闲时期消化这些计算,可以兼顾正则化效果和吞吐量。
首发原创文章，作者：王坚，如若转载，请注明出处：https://idctop.com/article/625493.html
