2

我正在处理一个太大而无法放入 RAM 的数据集。我目前正在尝试的解决方案是使用 numpy memmap 使用 Dataloader 一次加载一个样本/行。解决方案如下所示:

class MMDataset(torch.utils.data.Dataset):
    def __init__(self, path):
        self.file_path = path
        self.dataset_len = 44000000
        self.bytes_per_value = 32/8
        self.num_cols = 512
        self.num_rows = 1


    def __getitem__(self, index):


        x = np.memmap(self.file_path, dtype='float32', mode='r', shape=(
            self.num_rows, self.num_cols), offset=int(index*self.num_cols*self.bytes_per_value))

        return np.array(x)

    def __len__(self):
        return self.dataset_len



dataset = MMDataset('./data/emb.memmap')

data_loader = DataLoader(
    dataset,
    batch_size=4096,
    shuffle=True,
    num_workers=20
)

可用的 RAM 量大于 memmap 文件的大小时,数据加载速度很快。我得到大约60 批次/秒。但是,当可用 RAM 小于 memmap 文件的大小时,我得到大约 3 batches/second

我在尝试不同大小的 memmap 文件时发现了这一点。

为什么会这样?如果 Dataloader + memmap 将在可用 RAM < memmap 文件大小时节流,这将破坏解决方案的要点。

我观察到,当可用 RAM < memmap 文件大小时,磁盘 i/o 的读取速度为 500MB/s。这远高于加载一批 4096 个样本(接近 8MB/s)所需的理论读取量。

4

0 回答 0