h5py：使用分块存储向HDF5加载大数据时写入效率低下问题

核心问题定位 #

你的代码存在索引错误，这是写入速度极慢的主要原因：

• 你创建的数据集形状为(1024, 1024, 300)，前两个维度对应单个矩阵尺寸，第三个维度对应矩阵数量。
• 但代码中dset[ii]是取数据集第一个维度的第ii个切片（形状为(1024, 300)），与你加载的(1024,1024)矩阵不匹配，不仅会导致数据写入位置错误，还会触发HDF5大量随机分块修改——每个分块包含300个第三维度元素，每次修改其中一个都要读写整个分块，效率极低。

优化方案 #

1. 修正索引并调整分块大小 #

先修正写入索引，确保每个矩阵写入到第三维度的对应位置，同时调整分块大小匹配写入单元，避免分块部分修改：

import h5py
import numpy as np
from tqdm import tqdm

with h5py.File("FFT_Heights.h5", "w") as f:
    # 分块对应单个矩阵，写入时直接覆盖整个分块
    dset = f.create_dataset("chunked", (1024, 1024, 300), chunks=(1024, 1024, 1), dtype='complex128')
    for ii in tqdm(range(300)):
        # 修正索引：写入第三维度的第ii个切片
        dset[:, :, ii] = np.load(f'K field {ii}.npy').astype('complex128')

2. 批量读取写入（进一步提升效率） #

减少I/O操作次数，每次批量读取多个npy文件并一次性写入连续切片，大幅降低读写开销：

import h5py
import numpy as np
from tqdm import tqdm

batch_size = 16  # 根据内存情况调整，比如16、32或64
total_matrices = 300

with h5py.File("FFT_Heights.h5", "w") as f:
    # 分块大小匹配批量写入的第三维度规模
    dset = f.create_dataset("chunked", (1024, 1024, total_matrices), chunks=(1024, 1024, batch_size), dtype='complex128')
    
    for start_idx in tqdm(range(0, total_matrices, batch_size)):
        end_idx = min(start_idx + batch_size, total_matrices)
        # 批量加载npy文件
        batch_data = np.array([np.load(f'K field {ii}.npy').astype('complex128') for ii in range(start_idx, end_idx)])
        # 转置后写入对应切片
        dset[:, :, start_idx:end_idx] = batch_data.transpose(1, 2, 0)

3. 额外优化建议 #

• 关闭文件压缩：创建数据集时指定compression=None（默认值），避免压缩带来的CPU开销。
• 使用SSD存储：机械硬盘的随机读写性能远低于SSD，更换介质能显著提升写入速度。
• 关闭数据打乱：创建数据集时设置shuffle=False，减少写入时的额外操作。

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