Pytorch 高效快速加载大规模数据集
一、前言
最近遇到一个多模态学习任务,原始数据为HDF5 格式,遇到主要两个问题:一是数据量过大无法直接加载到内存,二是HDF5 是基于关键值索引,索引速度非常慢。在使用Pytorch 训练模型时,数据加载速度跟不上模型训练速度,导致GPU使用率低。阅读OLMO 框架关于数据集加载的源码,找到了一个基于二进制文件 memory-map 的解决方案。
HDF5 简介
HDF5(Hierarchical Data Format version 5)是一种用于存储和管理大量数据的文件格式,支持高效的多维数组存储,支持多层级数据模型,能够存储复杂的科学数据。HDF5 通常以groups 和 datasets 的形式组织。
import h5pydef print_dataset(name, obj):if isinstance(obj, h5py