Pytorch 高效快速加载大规模数据集

一、前言

最近遇到一个多模态学习任务，原始数据为HDF5 格式，遇到主要两个问题：一是数据量过大无法直接加载到内存，二是HDF5 是基于关键值索引，索引速度非常慢。在使用Pytorch 训练模型时，数据加载速度跟不上模型训练速度，导致GPU使用率低。阅读OLMO 框架关于数据集加载的源码，找到了一个基于二进制文件 memory-map 的解决方案。

HDF5 简介

HDF5（Hierarchical Data Format version 5）是一种用于存储和管理大量数据的文件格式，支持高效的多维数组存储，支持多层级数据模型，能够存储复杂的科学数据。HDF5 通常以groups 和 datasets 的形式组织。

import h5py

def print_dataset(name, obj):
    if isinstance(obj, h5py