如何处理多个大型.h5文件并创建PyTorch DataLoader

针对Ninapro DB2数据集拆分多个大HDF5文件、内存无法全量加载的问题，我们可以通过全局索引映射+按需读取的方式实现符合要求的DataLoader，确保每轮每个样本仅使用一次，且批次样本随机来自所有文件。

核心思路 #

1. 预先遍历所有HDF5文件，为每个样本建立全局唯一索引，记录该样本所属的文件路径、数据集类型（训练/测试）和本地索引，将这些信息保存为一个轻量级的索引文件。
2. 自定义PyTorch Dataset类，根据全局索引从对应文件中读取单条样本，避免全量加载数据。
3. 使用PyTorch DataLoader配合shuffle=True，实现每轮打乱全局索引，保证随机采样且无重复。

步骤1：生成全局样本索引映射 #

运行以下代码生成索引文件（仅需执行一次）：

import h5py
import pickle
from pathlib import Path

# 替换为你的8个HDF5文件路径列表
h5_file_paths = [
    "Sub1_5.h5",
    "Sub2_5.h5",
    "Sub3_5.h5",
    "Sub4_5.h5",
    "Sub5_5.h5",
    "Sub6_5.h5",
    "Sub7_5.h5",
    "Sub8_5.h5"
]

sample_index_map = []

for file_path in h5_file_paths:
    with h5py.File(file_path, 'r') as f:
        train_count = len(f['key1'])
        test_count = len(f['key3'])
        
        # 记录训练样本索引信息
        for local_idx in range(train_count):
            sample_index_map.append({
                'file_path': str(file_path),
                'type': 'train',
                'local_idx': local_idx
            })
        
        # 记录测试样本索引信息
        for local_idx in range(test_count):
            sample_index_map.append({
                'file_path': str(file_path),
                'type': 'test',
                'local_idx': local_idx
            })

# 保存索引映射到pickle文件
with open('ninapro_db2_global_index.pkl', 'wb') as f:
    pickle.dump(sample_index_map, f)

print(f"完成全局索引生成，总样本数: {len(sample_index_map)}")

步骤2：自定义PyTorch Dataset #

实现按需读取样本的Dataset类，内置文件句柄缓存以提升读取效率：

import torch
from torch.utils.data import Dataset
import h5py
import pickle

class NinaproDB2Dataset(Dataset):
    def __init__(self, index_file_path):
        # 加载全局索引映射
        with open(index_file_path, 'rb') as f:
            self.sample_index_map = pickle.load(f)
        # 缓存已打开的HDF5文件句柄，避免重复IO操作
        self._file_handles = {}

    def __len__(self):
        return len(self.sample_index_map)

    def __getitem__(self, global_idx):
        sample_info = self.sample_index_map[global_idx]
        file_path = sample_info['file_path']
        data_type = sample_info['type']
        local_idx = sample_info['local_idx']

        # 获取或打开目标文件句柄
        if file_path not in self._file_handles:
            self._file_handles[file_path] = h5py.File(file_path, 'r')
        h5_file = self._file_handles[file_path]

        # 读取单条样本和标签
        if data_type == 'train':
            data = h5_file['key1'][local_idx]
            label = h5_file['key2'][local_idx]
        else:
            data = h5_file['key3'][local_idx]
            label = h5_file['key4'][local_idx]

        # 转换为PyTorch张量，可根据模型需求调整维度和数据类型
        data_tensor = torch.tensor(data, dtype=torch.float32)
        label_tensor = torch.tensor(label, dtype=torch.long)

        return data_tensor, label_tensor

    def __del__(self):
        # 销毁时关闭所有打开的文件句柄
        for fh in self._file_handles.values():
            fh.close()

步骤3：构建DataLoader #

使用自定义Dataset创建满足要求的DataLoader：

from torch.utils.data import DataLoader

# 初始化数据集
dataset = NinaproDB2Dataset('ninapro_db2_global_index.pkl')

# 构建DataLoader，shuffle=True保证每轮打乱全局索引，实现无重复随机采样
dataloader = DataLoader(
    dataset,
    batch_size=512,
    shuffle=True,
    num_workers=4,  # 根据CPU核心数调整，多进程读取需保证HDF5文件为只读模式
    pin_memory=True  # GPU训练时开启，加速数据从CPU到GPU的传输
)

# 测试迭代示例
for batch_idx, (batch_data, batch_labels) in enumerate(dataloader):
    print(f"批次 {batch_idx+1}: 数据形状 {batch_data.shape}, 标签数量 {batch_labels.size(0)}")
    # 这里插入你的训练/验证逻辑
    if batch_idx == 9:  # 仅测试前10个批次
        break

关键注意事项 #

• 多进程兼容性：h5py支持多进程只读访问同一个文件，设置num_workers>0时每个worker会独立维护自己的文件句柄缓存，无需额外配置。
• 索引文件复用：生成的pickle文件仅几百KB，后续训练直接加载即可，无需重复生成。
• 数据拆分：若需要分离训练集和测试集，可在生成索引时将训练、测试样本的索引分别保存为两个文件，再创建对应的Dataset实例。
• 张量调整：根据模型输入要求，可修改__getitem__中数据张量的维度顺序（例如将(12,400)转为(400,12)）。

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