PyTorch在Linux上的数据加载优化方法

Linux下的PyTorch数据加载优化指南

一 基础配置与参数调优

• 使用SSD/NVMe存放数据，优先避免HDD带来的I/O瓶颈；将数据集放在高速盘可显著缩短读取时间。
• 合理设置num_workers：建议从CPU物理核心数起步，小模型或轻预处理可设为4–8，重预处理可逐步增大；不宜超过物理核心数，避免上下文切换与内存抖动。
• 启用pin_memory=True以使用锁页内存，加速CPU→GPU传输；结合non_blocking=True实现异步拷贝，重叠计算与传输。
• 提升预取能力：设置prefetch_factor=2（默认值），并使用persistent_workers=True（PyTorch≥1.7）避免每个epoch重建worker进程。
• 传输与计算重叠：在训练循环中，先启动下一轮数据预取，再执行当前轮计算，减少GPU空转。

二 存储与系统层优化

• 将热点数据放入**/dev/shm**（内存文件系统）以加速读取，适合中等规模数据集；注意其容量默认约为内存的一半，可按需调整：例如执行命令：sudo mount -o size=5128M -o remount /dev/shm。
• 图像解码加速：用pillow-simd替换Pillow，显著提升JPEG等解码性能（需按平台与依赖正确编译安装）。
• 系统层面：保持GPU驱动/CUDA/cuDNN/NCCL为较新稳定版本；按需调整文件描述符限制与内核网络参数；使用nvidia-smi、torch.autograd.profiler等工具持续观测GPU/CPU利用率与瓶颈。

三 缓存与预取策略

• 全内存预加载：当数据集能放入内存时，启动时一次性读入并常驻，后续直接从内存取样，I/O开销接近0（适合MNIST、CIFAR等小中型数据集）。
• 磁盘缓存：将预处理后的样本序列化为二进制文件（如joblib/pickle），首次构建、后续直接加载，适合大型固定数据集（如超分辨率、分割等）。
• 内存缓存：对频繁访问的样本使用LRU缓存，在内存中保留最近使用的若干图像，降低重复解码与增强成本。
• 异步预取器：实现或采用DataPrefetcher，用独立线程/流提前把下一批数据拷入GPU或在CPU侧排队，训练与加载并行，典型做法是用torch.cuda.Stream与后台线程预取。

四 分布式训练与多卡协同

• 多卡训练优先选择DistributedDataParallel（DDP），相较DataParallel具备更高效的梯度同步与通信；确保正确初始化进程组与NCCL通信库。
• 多卡场景下，常用经验是让每个进程的num_workers≈4×每卡GPU数，以喂饱各卡的GPU计算；同时开启pin_memory与prefetch_factor提升端到端吞吐。
• 结合**AMP（torch.cuda.amp）**降低显存占用并提升吞吐，与数据加载优化叠加效果更佳。

五 快速检查清单与示例配置

• 检查清单
  • 存储：数据在SSD/NVMe；热点数据可尝试**/dev/shm**；必要时做磁盘/内存缓存。
  • DataLoader：num_workers≈CPU物理核数；pin_memory=True；prefetch_factor=2；persistent_workers=True。
  • 传输：使用non_blocking=True与DataPrefetcher重叠计算与传输。
  • 监控：用nvidia-smi观察GPU利用率，目标长期不低于70%；用torch.utils.bottleneck定位瓶颈。
• 示例配置（单节点多卡，示意）
  • num_workers：设为CPU物理核心数（或每卡4并随核数微调）
  • pin_memory：True
  • prefetch_factor：2
  • persistent_workers：True
  • 传输：tensor.to(device, non_blocking=True)
  • 训练循环：启动下一轮预取→当前轮前向/反向→同步与步进
  • 存储：优先SSD/NVMe；热点数据可放**/dev/shm**；固定数据集建议磁盘缓存
  • 监控：nvidia-smi dmon；必要时运行python -m torch.utils.bottleneck定位数据环节瓶颈