如何通过H5py加速机器学习场景下300MB级图像样本读取？

H5py本身的优化设置 #

1. 优化分块策略 #

你当前的分块大小远超磁盘IO最优范围，导致读取效率低下。需调整分块大小匹配磁盘的最优IO块（通常64MB-128MB），同时对齐你的读取模式：

• 若按单个样本读取，可将单样本拆分为多个子区域分块，确保每个分块大小落在64-128MB区间。比如假设图像维度为(H, W, 3)（RGB），300MB的样本可拆分为(H//4, W, 3)，单块约75MB：

  # 假设resolution为(H, W, 3)
  chunk_h = resolution[0] // 4
  data_group.create_dataset(k_name, data=knth, chunks=(chunk_h, resolution[1], resolution[2]))
  

• 若为批量读取多样本，可设置分块为(batch_size, *sub_resolution)，保证单块大小在64-128MB区间，避免跨块读取的额外开销。

2. 调整读取缓存参数 #

h5py默认缓存较小，增大缓存可减少重复读取（如训练多epoch）时的磁盘IO次数。打开文件时设置缓存总大小和缓存块数量：

with h5py.File("src.h5", "r", rdcc_nbytes=1024*1024*1024, rdcc_nslots=100) as f:
    sample = load_data(f)

这里设置1GB缓存，可根据内存余量调整（如内存充足可设为2GB）。

3. 禁用压缩（若启用过） #

如果创建数据集时使用了compression参数（如gzip），读取时的解压计算会拖慢速度。写入时关闭压缩，优先保证读取性能：

# 创建数据集时不指定compression参数
data_group.create_dataset(k_name, data=knth)

4. 并行读取样本 #

利用多进程并行加载多个样本，充分发挥磁盘的并行IO能力：

from multiprocessing import Pool

def load_single_sample(sample_name):
    with h5py.File("src.h5", "r", rdcc_nbytes=512*1024*1024) as f:
        return f[sample_name][...]

# 并行加载5个样本
with Pool(5) as p:
    samples = p.map(load_single_sample, ["sample1", "sample2", "sample3", "sample4", "sample5"])

H5py的替代方案 #

若h5py的优化无法满足需求，可尝试以下专为机器学习场景优化的存储方案：

1. Zarr #

Zarr与h5py API高度兼容，天生支持分块存储和并行读取，大样本批量读取性能更优，且无需依赖HDF5库：

import zarr

# 写入示例（单样本数据集）
store = zarr.DirectoryStore("data.zarr")
root = zarr.group(store=store)
root.create_dataset(k_name, data=knth, chunks=(chunk_h, resolution[1], resolution[2]))

# 读取示例
with zarr.open("data.zarr", "r") as f:
    sample = f[k_name][...]

2. LMDB #

LMDB是内存映射的键值数据库，读取速度极快，适合单样本较大的场景，首次加载后几乎无磁盘IO开销：

import lmdb
import numpy as np

# 写入示例
env = lmdb.open("data.lmdb", map_size=200*300*1024*1024)  # 预分配60GB空间
with env.begin(write=True) as txn:
    txn.put(k_name.encode(), knth.tobytes())

# 读取示例
env = lmdb.open("data.lmdb", readonly=True)
with env.begin() as txn:
    data_bytes = txn.get(k_name.encode())
    sample = np.frombuffer(data_bytes, dtype=np.uint8).reshape(resolution)

3. Webdataset #

将每个样本保存为单独文件并打包成tar包，Webdataset可高效批量读取tar中的样本，支持并行加载和预取，适配分布式训练：

# 命令行打包样本为tar
tar -cf dataset.tar sample1.png sample2.png ...

import webdataset as wds

dataset = wds.WebDataset("dataset.tar").decode("pil").to_tuple("png")
for sample in dataset:
    # 处理样本逻辑
    pass

4. TFRecord/PyTorch 自定义Dataset #

• 若使用TensorFlow，可将H5数据转换为TFRecord，它支持并行读取、预取和流水线式数据增强；
• 若使用PyTorch，可将所有样本存储为连续二进制文件，结合numpy.memmap实现零拷贝读取，自定义Dataset加载。

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