Python中读取张量文件的最快方法 适配PyTorch训练需求

你当前读取速度慢的核心原因是h5py压缩文件的随机读取需要额外解压开销，且批量打包的存储结构会提升单样本寻址成本。以下是无需全量预加载到内存的高读取速度方案，按优先级排序：

• 单独未压缩NumPy .npy格式（最优选择）
  .npy是内存连续的原生二进制格式，无编解码开销，支持内存映射读取，完全无需加载全量数据到内存。单个张量对应一个.npy文件，读取时直接调用代码：

  import numpy as np
  import torch
  # mmap_mode='r' 开启内存映射，仅加载访问的分片到内存
  tensor = torch.from_numpy(np.load("sample_xxx.npy", mmap_mode="r"))
  

  该方案单样本随机读取速度是当前h5压缩方案的3~5倍，适配PyTorch DataLoader多进程读取逻辑，调试成本最低。
• LMDB键值数据库（适合小文件管理场景）
  如果需要避免生成大量单样本小文件占用过多文件系统inode，可选择LMDB做打包存储。将所有张量序列化后存入LMDB，键为样本ID，值为序列化的二进制数据，LMDB本身基于内存映射实现，随机读取性能接近单独.npy文件，比h5压缩方案快2~4倍，同样无需全量加载数据到内存。
• 未压缩PyTorch .pt/.pth格式
  不需要跨框架兼容的场景下可使用，存储时关闭zip序列化压缩，读取时直接调用torch.load("sample_xxx.pt", map_location="cpu")即可，速度略低于.npy格式，无需做NumPy到Torch张量的转换。

避坑注意事项 #

• 所有存储方案均不要开启压缩，压缩/解压开销会大幅抵消二进制读取的性能优势，优先用磁盘空间换取读取速度。
• 配合PyTorch DataLoader使用时开启多进程（设置合理的num_workers参数），可完全发挥上述存储格式的读取性能。
• 若必须保留h5打包存储结构，可将h5文件切换为未压缩模式，读取速度也可提升2倍左右，但性能仍弱于前两种方案。

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