快速使用 HDF5 存储数据

遇到要把分散的训练集合（雷达数据、标签等）合在一个文件的的需求，然后看上了 HDF5 这种存储方式。

HDF5 还支持数据压缩，并且读取数据时会自动解压。

以下代码实例使用库 h5py。

import h5py

with h5py.File(output_file, 'w') as f:
    all_radar_data = []
    all_class_data = []
    all_box_data = []
    for file in common_files:
        ......

    all_radar_data = np.array(all_radar_data)
    all_class_data = np.array(all_class_data)
    all_box_data = np.array(all_box_data)

    f.create_dataset('raw_data', data=all_radar_data, compression='gzip', compression_opts=8)
    f.create_dataset('class', data=all_class_data)
    f.create_dataset('box', data=all_box_data)

关于压缩，compression='gzip', compression_opts=8 就启用了数据压缩。compression 设置为 gzip 就行，compression_opts 值越高压缩力度越大，最高至 9。实测 9 压缩慢到令人发指，读取速度则是各个压缩等级都一样。

还有另一种先立靶子再添加数据的存储方式：

dset = f.create_dataset(
            '3dfft_data',
            (num_files, 16, 244, 244),
            compression='gzip',
            compression_opts=9,
            dtype='float32',
        )

...

dset[start_index : start_index + len(radar_data_list)] = radar_data_concat

这个方式需要先写定数据的维度大小。h5py 还提供了一种 vlen 的方法允许不定长的数据存入，具体可见文档。

文件的读取非常简单：

with h5py.File(file_name, 'r') as f:
    raw_data = f['raw_data'][:]

实际用下来，使用 compression_opts=8 压缩率大概在 60％，压缩 800MB 的数据消耗 15 分钟。读取耗时 10s 以内。

加快读取速度

若将 HDF5 文件封装为 torch.utils.data.DataLoader，可能会发现加载 batch 奇慢无比，且相当消耗 CPU 资源。

写入 HDF5 文件时设置 chunk

HDF5 文件以 chunk 为单位进行数据存取。要读出数据，就要读出数据所在的整个 chunk。

h5py 的自动 chunk 大小极有可能不符合实际需求。最好手动设置为一次数据的大小。

dset = f.create_dataset(
            '3dfft_data',
            (num_files, 16, 244, 244),
            maxshape=(None, 16, 244, 244),
            chunks=(1, 16, 244, 244), # 手动设置 chunk 大小为一次数据的大小
            compression='gzip',
            compression_opts=9,
            dtype='float32',
        )

...

读取 HDF5 文件时设置缓存

Chunk cache 是 h5py 默认启用的特性，用于在读取 HDF5 文件时将某些 chunk 缓存在内存，加速读取。

但 h5py 的默认缓存大小只有 1MB。可以通过 rdcc 系列参数进行设置。

f = h5.File(
    File_Name_HDF5,
    'r',
    rdcc_nbytes=1024**2 * 4000,  # 缓存总大小，单位字节。应设置为 chunk 的整数倍
    rdcc_w0=1,  # 缓存淘汰策略。值越接近 0，优先淘汰最近最少使用的 chunk。值越接近 1，优先淘汰已完全读取或写入的块
    rdcc_nslots=1e7,  # 定位缓存 chunk 位置的坐标长度。推荐为缓存 chunk 数的 10 倍，为达最佳性能需要 100 倍。默认为 521
)

参考来源

• https://stackoverflow.com/questions/48385256/optimal-hdf5-dataset-chunk-shape-for-reading-rows