比较深度学习中数据加载方法的效率

Question

我需要加载时间序列数据集来训练网络。当我从原始数据中提取这些.npy 文件时，由于内存问题，数据集被分成许多块train_x_0.npy、train_x_1.npy、...、train_x_40.npy（41 个块）。但是，它们的大小太大（大约 1000 GB），我无法将所有内容加载到 RAM 中。我一直在考虑两种方法来解决这个问题。

1. 使用np.load() 和参数mmap_mode='r+' 加载数据块。内存映射的块存储在 Python 列表self.data 中。在PytorchDataset类的__getitem__(self, idx)方法中，我将idx转换为chunk_idx和sample_idx，然后通过self.data[chunk_idx][sample_idx]获取样本。
2. 再次从原始数据中提取.npy文件，并逐个样本保存数据，即一个.npy文件现在是一个样本，而不是一个数据块。在__getitem__(self, idx) 方法中，我将通过使用np.load(sample_path) 加载它来获取一个样本。

假设 Pytorch DataLoader 将用于遍历所有样本，那么哪种方法会更快？

如果您对提取原始数据或加载.npy 文件有其他建议，请分享您的意见。

Answer 1

两种建议的方法都将受到文件系统 IO 的限制，因为每个样本都将按需从磁盘加载（内存映射不会加速实际加载，一旦请求给定补丁）。

特别是当您计划训练多个 epoch 时，您可以通过一次加载一个原始块 train_x_0.npy、train_x_1.npy 等（或尽可能多的 RAM 中）来实现强大的加速并在切换到下一个之前在这个块上训练多个 epoch。

为此，您需要控制dataloader 请求的样本索引。为此，您可以定义一个采样器，该采样器传递相应缓存数据块中可用的示例索引。在伪代码中，一次缓存一个块时，您的训练循环可能看起来像这样：

from yourproject import Dataset
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler

dataset = Dataset(train_data_path, ...)
for chunk_idx in range(num_chunks):
  dataset.cache_chunk(chunk_idx)
  chunk_sample_inds = dataset.get_chunk_sample_inds(chunk_idx)
  chunk_sampler = SubsetRandomSampler(chunk_sample_inds)
  chunk_loader = DataLoader(dataset=dataset, sampler=chunk_sampler)
  for chunk_epoch in range(num_chunk_epoch):
    for sample_idx, sample in enumerate(chunk_loader):
       output = model(sample)

在此，您的Dataset 班级需要照顾

• 缓存（加载到 RAM）指定块，给定块 idx（由 cache_chunk 方法指示）
• 返回给定块 idx 的有效样本索引列表（由 get_chunk_sample_inds 方法指示）

如果您使用快速 GPU（这通常受到在 RAM 和 VRAM 之间来回移动数据的限制，即使对于 RAM 缓存的数据也是如此），您可以预期使用这种方法可以提高几个数量级的速度（而不是尝试为每个样本从 HDD 填充 VRAM）。