加载多个大型 H5PY 文件时的 I/O 问题 (Pytorch)

Question

我遇到了问题！

最近遇到一个 I/O 问题。目标和输入数据存储在 h5py 文件中。每个目标文件为 2.6GB，而每个输入文件为 10.2GB。我总共有 5 个输入数据集和 5 个目标数据集。

我为每个 h5py 文件创建了一个自定义数据集函数，然后使用data.ConcatDataset 类链接所有数据集。自定义数据集函数为：

class MydataSet(Dataset):
def __init__(self, indx=1, root_path='./xxx', tar_size=128, data_aug=True, train=True):
    self.train = train
    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 = h5py.File(self.in_file, 'r', driver='core')
    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())
    # h5f_n.close()
    # h5f_c.close()

    self.tar_size = tar_size
    self.data_aug = data_aug

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

def __del__(self):
    self.h5f_n.close()
    self.h5f_c.close()

def __getitem__(self, index):
    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])

然后我使用dataset_train = data.ConcatDataset([MydataSet(indx=index, train=True) for index in range(1, 6)]) 进行训练。当只使用 2-3 个 h5py 文件时，I/O 速度正常，一切正常。但是，当使用 5 个文件时，训练速度逐渐降低（5 次迭代/秒到 1 次迭代/秒）。我换了num_worker，问题依旧存在。

谁能给我一个解决方案？我应该将几个 h5py 文件合并成一个更大的文件吗？还是其他方法？提前致谢！

Answer 1

提高性能需要时间基准。为此，您需要识别潜在的瓶颈和相关场景。您说“使用 2-3 个文件，I/O 速度正常”和“使用 5 个文件时，训练速度逐渐降低”。那么，您的性能问题是 I/O 速度还是训练速度？或者你知道吗？如果不知道，则需要针对 2 个场景分别隔离比较 I/O 性能和训练性能。
换句话说，要（仅）测量 I/O 性能，您需要运行以下测试：

1. 是时候读取和连接 2-3 个文件了，
2. 是时候读取和连接 5 个文件了，
3. 将5个文件复制成1个，并从合并的文件中读取时间，
4. 或者，将 5 个文件链接到 1 个文件和时间。

要衡量训练速度（仅），您需要比较以下测试的性能：

• 合并 2-3 个文件，然后从合并的文件中读取和训练。
• 合并所有 5 个文件，然后从合并的文件中读取和训练。
• 或者，将 5 个文件链接到 1 个文件，然后从链接文件中读取和训练。

正如我在评论中指出的，如果所有数据集都位于根级别并且所有数据集名称都是唯一的，那么将多个 HDF5 文件合并（或链接）为一个文件很容易。我添加了外部链接方法，因为它可能提供相同的性能，而无需复制大型数据文件。

下面是显示这两种方法的代码。在fnames 列表中替换您的文件名，它应该可以运行了。如果您的数据集名称不是唯一的，则需要创建唯一的名称，并在 h5fr.copy() 中分配——像这样：h5fr.copy(h5fr[ds],h5fw,'unique_dataset_name')

合并或链接文件的代码：
（酌情评论/取消评论行）

import h5py
fnames = ['file_1.h5','file_2.h5','file_3.h5']
# consider changing filename to 'linked_' when using links:
with h5py.File(f'merge_{len(fnames)}.h5','w') as h5fw:      
    for fname in fnames:
        with h5py.File(fname,'r') as h5fr:
            for ds in h5fr.keys():
                # To copy datasets into 1 file use:
                h5fr.copy(h5fr[ds],h5fw)
                # to link datasets to 1 file use:
                # h5fw[ds] = h5py.ExternalLink(fname,ds)