Ubuntu如何优化PyTorch的I/O性能

Ubuntu下优化PyTorch的I/O性能

一 存储与文件系统

• 使用NVMe SSD作为数据盘，顺序读取带宽可达约7 GB/s，远高于机械盘，能显著缩短数据读取与解压等待时间。
• 将热点数据放入内存文件系统**/dev/shm**（tmpfs）以减小I/O延迟：
  • 查看容量：df -h /dev/shm
  • 调整大小（示例为4GB）：在/etc/fstab添加行：tmpfs /dev/shm tmpfs defaults,size=4096M 0 0
  • 重新挂载：sudo umount /dev/shm && sudo mount /dev/shm
  • 再次检查：df -h /dev/shm
    以上措施能明显降低数据加载的磁盘等待，特别适合数据集可放入内存的场景。

二 DataLoader与数据预取

• 并行加载：将DataLoader的num_workers设为大于0并合理调优（通常设为CPU物理核心数或略高），以隐藏磁盘与解码延迟。
• 固定内存：开启pin_memory=True，并使用non_blocking=True进行主机到GPU的异步拷贝，提升传输与GPU利用率。
• 预取与顺序读：尽量顺序读取、批量读取，减少随机小I/O；在Dataset.__getitem__中避免高开销的重复操作。
• 高效图像解码：使用Pillow-SIMD、turbojpeg或jpeg4py替代纯Python/PIL解码路径，降低CPU解码瓶颈。
• 典型示例：

  from torch.utils.data import DataLoader
  
  loader = DataLoader(
      dataset,
      batch_size=256,
      shuffle=True,
      num_workers=8,          # 依据CPU核心数调优
      pin_memory=True,        # 启用固定内存
      persistent_workers=True, # 长时训练减少进程反复创建开销
  )
  
  for batch in loader:
      x, y = batch
      x = x.to(device, non_blocking=True)  # 异步拷贝到GPU
  

  上述配置可显著减少I/O等待并提升GPU利用率。

三 数据格式与管线

• 采用高性能存储格式：将大规模数据转换为HDF5或LMDB，以获得更好的顺序读取与随机访问性能，减少Python层开销。
• 使用GPU加速数据管线：引入NVIDIA DALI进行图像解码与增强，降低CPU解码压力，特别适合高分辨率图像与多卡训练。
• 示例（HDF5读取思路）：

  import h5py, torch
  def collate(batch):
      xs, ys = [], []
      for x, y in batch:
          xs.append(torch.from_numpy(x))
          ys.append(y)
      return torch.stack(xs), torch.tensor(ys)
  
  with h5py.File('data.h5', 'r') as f:
      # 假设数据按dataset_name组织
      dset_x = f['images']
      dset_y = f['labels']
      # 结合DataLoader与collate_fn使用
  

  这些做法能在数据规模大、解码复杂时有效缓解I/O瓶颈。

四 系统与内核参数

• 提升文件描述符限制：在/etc/security/limits.conf增加如* soft nofile 65536、* hard nofile 65536，并重启会话，避免“Too many open files”。
• I/O调度与调度器：对NVMe设备优先使用none/mq-deadline等低开销调度器（视内核与发行版而定），减少调度开销。
• 其他内核参数：可按需调整vm.swappiness、vm.dirty_ratio等，减少抖动；在容器/虚拟化环境中确保存储直通与I/O配额合理。
• 环境一致性：保持GPU驱动、CUDA、cuDNN、NCCL版本匹配并更新，避免因兼容性问题导致隐式性能退化。
  这些系统级优化能为高并发数据加载与高速存储提供稳定的底层支持。

五 监控与瓶颈定位

• 实时监控：
  • GPU：nvidia-smi dmon观察GPU利用率、显存与功耗；
  • 存储：iostat -x 1查看磁盘util、await、svctm等指标；
  • 系统：htop/top观察CPU与内存压力。
• 性能剖析：
  • 使用torch.profiler定位数据加载、传输与计算各阶段耗时；
  • 结合TensorBoard可视化，验证优化是否有效。
• 迭代验证：每次只调整一个变量（如num_workers、是否使用pin_memory、是否迁移到**/dev/shm**），用上述工具对比前后差异，确保收益可复现。
  通过系统化的监控与剖析，能快速识别I/O瓶颈并量化优化成效。