PyTorch DataLoader如何优化大体积npz数据集的批次读取效率

核心传参逻辑调整 #

• 首先将预先按等长切分好的.npz文件块列表存为类成员变量self.data_blocks，每个元素对应单个缓存块包含的所有样本路径；同步在类初始化阶段声明两个状态变量，替代无参调用load_data的逻辑：
  • self.cached_block_idx = -1：标记当前内存中缓存对应的数据块索引，初始状态为未加载任何块
  • self.cached_samples = None：存储内存中实际加载的批量样本张量
• 修正load_data方法的定义规则：将入参设置为块索引block_idx，方法内部逻辑为根据传入的索引从self.data_blocks取出对应路径列表，批量将块内所有.npz文件加载为连续内存张量，赋值给self.cached_samples，同时更新self.cached_block_idx = block_idx。
  注意：不要在类初始化阶段直接调用load_data，加载动作完全由样本访问流程动态触发，从根源上解决无参调用无法匹配索引的问题。

内存缓存调度流程打通 #

• 改造__getitem__方法的样本获取逻辑，替代原逐样本磁盘读取流程：
  1. 传入单样本全局索引idx后，根据预切分的等长块规则计算当前样本所属的块索引：block_idx = idx // self.block_size，其中self.block_size就是预设的缓存容量（即num_gpus×num_workers×batch_size量级）
  2. 比对计算出的block_idx和当前self.cached_block_idx：如果两者相等，直接根据块内偏移量block_offset = idx % self.block_size从self.cached_samples中读取样本返回，全程仅做内存访问，无磁盘IO开销
  3. 如果两者不相等，触发缓存替换流程：调用self.load_data(block_idx)加载新块数据，旧块数据会被直接覆盖回收内存，不会造成持续的内存占用上涨
• 该流程完全匹配设计目标：一个缓存块加载完成后，块范围内所有被请求的样本（覆盖所有GPU、所有DataLoader worker的访问请求）都会直接走内存读取，直到块内所有样本被消费完毕，才会触发下一次批量磁盘读取，彻底消除逐样本随机读的性能瓶颈。

多卡分布式适配注意事项 #

• 因为搭配DistributedSampler做样本分发，只要预切分的data_blocks是基于全局样本列表生成的，多进程DataLoader的每个worker独立内存空间的缓存不会出现额外的冗余占用，只要控制单块缓存大小不超出单机内存上限即可
• 配合pin_memory=True参数时，缓存加载的样本直接以连续张量形式存入锁页内存，可直接复用原有CPU到GPU的传输加速逻辑，不需要额外适配
• 批量加载块数据时，要避免逐样本单独打开.npz文件读取，尽量将块内样本的文件句柄批量打开、拼接为连续内存数组后再转成张量，进一步降低文件系统寻道的额外开销。

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