我正在处理一个太大而无法放入 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)所需的理论读取量。