关键在于优化数据加载以避免拖慢GPU,核心方法包括预加载、并行化、零拷贝和缓存复用;通过设置num_workers、pin_memory、persistent_workers,改用torchvision.io.read_image、预处理存储为.pt/.lmdb格式、IterableDataset分片流式加载,并结合性能监控定位瓶颈。

关键不是换框架,而是让数据加载不拖GPU后腿。核心思路是:预加载、并行化、零拷贝、缓存复用。
用DataLoader的多进程和异步加载
PyTorch默认单线程读数据,极易成为瓶颈。设置num_workers > 0启用子进程预取,配合pin_memory=True把数据提前搬进GPU可直接访问的锁页内存。
- num_workers建议设为CPU逻辑核数的1–2倍(如8核机器试4–6),太多反而因进程调度开销变慢
- 确保persistent_workers=True,避免每个epoch重建进程,减少初始化延迟
- 若训练中报“OSError: too many open files”,调高系统ulimit -n,或适当减小num_workers
避免运行时重复解码与变换
图像每次读取都解码(JPEG→Tensor)、再做ToTensor、Normalize,非常耗时。优先在数据准备阶段完成静态处理。
- 用torchvision.io.read_image替代PIL.Image.open,速度提升2–3倍,且原生支持uint8张量输出
- 对固定尺寸、无需随机裁剪的数据集,提前将图像转成.pt或.lmdb格式,加载时直接mmap读取,跳过解码
- 自定义Dataset中,把__getitem__里耗时操作(如OpenCV滤波)移到__init__阶段预计算并缓存
用IterableDataset + 流式分片应对超大数据集
当数据远大于内存(如千万级图像),传统random shuffle会卡死。改用分片+流式迭代更可控。
立即学习“Python免费学习笔记(深入)”;
- 把数据按10k–100k样本/份切分成多个shard文件(如train_0001.pt, train_0002.pt)
- 用IterableDataset按需加载当前shard,shuffle在shard内进行,epoch末打乱shard顺序
- 配合torch.utils.data.ChainDataset串联多个shard,避免一次性加载全部索引
监控瓶颈,别靠猜
用torch.utils.benchmark或简单time.time()测单次next(iter(dataloader))耗时,对比GPU空转时间(如torch.cuda.synchronize()前后打点)。
- 若数据加载耗时 > GPU计算耗时的1.5倍,说明I/O严重拖累
- nvidia-smi观察GPU利用率长期低于60%,同时CPU使用率满载 → 典型数据加载瓶颈
- 用py-spy record -p
采样,看Python线程是否卡在PIL、cv2或pickle.load上
基本上就这些。优化效果往往立竿见影——从每batch 200ms降到30ms很常见。不复杂但容易忽略。










