PyTorch DataLoader调用np.load在NVME SSD上读取数据极慢问题求助

解决方案 #

• 修复代码显性错误
  1. 补全路径拼接逻辑：os.listdir仅返回文件名，加载时需拼接完整路径：path = os.path.join("/home/bla/data", self.paths[idx])
  2. 修正初始化笔误：__init__方法中self.paths = np.array(self.paths)需改为self.paths = np.array(paths)，避免未赋值先调用的隐性异常
• 调整DataLoader参数配置
  1. 新增persistent_workers=True参数（PyTorch 1.9及以上版本支持），避免每轮迭代销毁重建worker进程，直接解决每N次迭代卡顿的问题
  2. 把prefetch_factor调整为4~8，拉长worker的数据预取队列长度，避免预取队列空转导致GPU等待数据
  3. 若使用PyTorch 2.0及以上版本，可新增pin_memory_device="cuda"参数，进一步降低CPU到GPU的数据拷贝开销
• 解决多进程线程竞争问题
  在Python脚本最开头加入以下环境变量配置，禁用numpy内部多线程，避免多worker场景下线程争抢CPU资源：

  import os
  os.environ["OMP_NUM_THREADS"] = "1"
  os.environ["OPENBLAS_NUM_THREADS"] = "1"
  os.environ["MKL_NUM_THREADS"] = "1"
  os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
  os.environ["NUMEXPR_NUM_THREADS"] = "1"
  

• 优化存储读取效率
  1. 检查NVME SSD的IO调度器：执行cat /sys/block/nvme1n1/queue/scheduler，确认使用none或mq-deadline调度器，若为机械盘适配的cfq调度器，可执行echo none > /sys/block/nvme1n1/queue/scheduler修改
  2. 降低小文件随机读开销：将所有npy文件打包为WebDataset格式的tar包，顺序读取可跑满SSD的2GB/s带宽
  3. 若保留独立npy文件存储，可针对ext4格式磁盘执行e4defrag /dev/nvme1n1整理碎片，提升随机读性能
• 快速验证方案
  先将shuffle=True改为shuffle=False测试读取速度，若速度恢复正常，说明随机读取是核心瓶颈，可通过增大预取队列、分块shuffle的方式进一步优化

内容的提问来源于stack exchange，提问作者JoSauderGH