请求将TensorFlow的HDF5数据生成器改写为PyTorch DataLoader实现

对应你提供的TensorFlow数据生成器，PyTorch中我们通常用Dataset和DataLoader实现高效分批数据加载，这两个是PyTorch数据加载的标准组件，原生支持shuffle、多进程加载、预取等功能，完全替代自定义生成器逻辑。

实现代码 #

import torch
from torch.utils.data import Dataset, DataLoader
import h5py
import numpy as np

class CustomHDF5Dataset(Dataset):
    def __init__(self, hdf5_path, name_x, name_y):
        # 以只读模式打开HDF5文件
        self.hdf5_file = h5py.File(hdf5_path, "r")
        # 保存数据集引用（不会立即加载全量数据到内存）
        self.x = self.hdf5_file[name_x]
        self.y = self.hdf5_file[name_y]
        # 记录总样本数
        self.num_samples = self.x.shape[0]

    def __len__(self):
        # 返回数据集总样本数
        return self.num_samples

    def __getitem__(self, idx):
        # 根据索引读取单个样本，转换为PyTorch张量
        x_sample = torch.from_numpy(np.array(self.x[idx], dtype=np.float32))
        y_sample = torch.from_numpy(np.array(self.y[idx], dtype=np.float32))
        return x_sample, y_sample

# 封装DataLoader创建逻辑
def get_dataloader(hdf5_path, batch_size, name_x, name_y, shuffle=True, num_workers=2):
    dataset = CustomHDF5Dataset(hdf5_path, name_x, name_y)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        prefetch_factor=2,  # 预取批次数量，对应TensorFlow的prefetch(1)
        persistent_workers=True  # 保持进程存活，提升加载效率
    )

关键功能对应说明 #

• 内存优化：HDF5数据集仅保存引用，只有当__getitem__被调用时才读取对应索引的样本，避免一次性加载全量数据。
• 替代原生成器逻辑：
  • 样本打乱：DataLoader的shuffle=True参数自动实现每个epoch重新打乱样本顺序，替代原代码中手动维护sample_nums和pointer的逻辑。
  • 分批拼接：batch_size参数直接指定批次大小，DataLoader自动完成样本拼接，无需手动收集单样本再转数组。
  • 预取加速：prefetch_factor控制预取的批次数量，实现数据加载与模型训练的并行，和TensorFlow的prefetch功能一致。

训练循环示例 #

# 创建训练数据加载器
train_loader = get_dataloader("train_data.h5", batch_size=32, name_x="input", name_y="label")

# 迭代训练
for epoch in range(10):
    for batch_x, batch_y in train_loader:
        # 在此执行模型前向传播、损失计算、反向传播等操作
        print(f"当前批次形状：x={batch_x.shape}, y={batch_y.shape}")

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