PyTorch加载多个大型h5py文件出现I/O性能下降问题该如何解决？

问题根因 #

你遇到的训练速度逐步下降的核心原因是初始化MydataSet时对输入h5文件使用了driver='core'参数，这个参数会将整个h5文件完整加载到系统内存中。按你的数据规模计算，5组数据集全部加载到内存需要占用约(10.2GB + 2.6GB)*5 = 64GB内存，当内存占用超过系统可用内存阈值后，系统会触发内存交换（swap），将部分内存数据转储到磁盘，导致I/O速度骤降，训练速度自然越来越慢。

1. 优先修复h5文件加载逻辑 #

• 去掉h5py.File的driver='core'参数，改为默认的磁盘读写模式，避免初始化时一次性加载全量文件到内存。
• 延迟加载h5文件句柄：不要在__init__里直接打开h5文件，而是在第一次调用__getitem__时再打开，避免多进程加载时出现句柄复制错误的问题，修改参考代码如下：

class MydataSet(Dataset):
    def __init__(self, indx=1, root_path='./xxx', tar_size=128, data_aug=True, train=True):
        self.train = train
        self.indx = indx
        self.root_path = root_path
        self.tar_size = tar_size
        self.data_aug = data_aug
        # 初始化时只存路径，不打开文件
        if self.train:
            self.in_file = pth.join(root_path, 'train', 'train_noisy_%d.h5' % indx)
            self.tar_file = pth.join(root_path, 'train', 'train_clean_%d.h5' % indx)
        else:
            self.in_file = pth.join(root_path, 'test', 'test_noisy.h5')
            self.tar_file = pth.join(root_path, 'test', 'test_clean.h5')
        self.h5f_n = None
        self.h5f_c = None
        self.keys_n = None
        self.keys_c = None

    def _init_h5(self):
        # 第一次取数据时再初始化文件句柄
        if self.h5f_n is None:
            self.h5f_n = h5py.File(self.in_file, 'r')
            self.h5f_c = h5py.File(self.tar_file, 'r')
            self.keys_n = list(self.h5f_n.keys())
            self.keys_c = list(self.h5f_c.keys())

    def __len__(self):
        self._init_h5()
        return len(self.keys_n)

    def __del__(self):
        if self.h5f_n is not None:
            self.h5f_n.close()
            self.h5f_c.close()

    def __getitem__(self, index):
        self._init_h5()
        keyn = self.keys_n[index]
        keyc = self.keys_c[index]
        datan = np.array(self.h5f_n[keyn])
        datac = np.array(self.h5f_c[keyc])
        datan_tensor = torch.from_numpy(datan).unsqueeze(0)
        datac_tensor = torch.from_numpy(datac)
        if self.data_aug and np.random.randint(2, size=1)[0] == 1: # horizontal flip
            datan_tensor = torch.flip(datan_tensor,dims=[2]) # c h w
            datac_tensor = torch.flip(datac_tensor,dims=[2])
        return datan_tensor, datac_tensor

2. 可选优化：合并h5文件 #

合并为单个文件不是必须的，但可以进一步减少磁盘随机寻址开销：把所有训练数据合并到同一个h5文件中，避免同时打开多个文件句柄带来的系统开销，也方便后续统一做数据切块、预读取优化。

3. 其他I/O优化方案 #

• 开启DataLoader的pin_memory=True参数，加快数据从内存到GPU的传输速度；如果使用多进程加载，设置persistent_workers=True避免每个epoch都重新初始化worker进程，减少重复打开h5文件的开销。
• 如果内存足够，可对高频访问的数据做缓存：在__getitem__里将读取过的张量存入内存缓存，下次访问时直接返回，减少重复读盘的次数。
• 对h5文件做分块压缩存储，写入h5时指定chunks=True和合适的压缩参数，减少文件体积，降低读盘的I/O带宽占用。

内容的提问来源于stack exchange，提问作者stander Qiu