存储和读取海量3D数据集的空间高效方式？

问题描述

我正在尝试对顺序数据进行神经网络训练。我的数据集将包含360万个训练示例。每个示例将是30 x 32 ndarray(在30天内观察到32个特征)。

我的问题是，写入和读取此数据最节省空间的方式是什么？

本质上它将具有(3.6m, 30, 32)和np.save()的形状，看起来很方便，但是我不能将整个内容保存在内存中，所以我不能真正使用np.save()保存它(或者使用np.load()重新加载它)。CSV也不起作用，因为我的数据有3个维度。

我的创建计划是成批处理条目，并将它们追加到某个文件中，这样我就可以在执行过程中释放内存。

最终，我将使用该数据文件作为PyTorch IterableDataset的输入，因此它必须是可以一次加载一行的文件(类似于.txt文件，但我希望有更好的方法来保存该数据，使其更符合表格的三维特性)。欢迎提出任何意见！

---

解决方案

因为您计划使用可迭代的数据集，所以不应该需要随机访问(IterableDataset不支持随机采样器)。在这种情况下，为什么不将所有内容都写到一个二进制文件中并对其进行迭代呢？我发现在实践中，这通常比其他解决方案要快得多。这应该比另存为文本文件快得多，因为您避免了将文本转换为数字的开销。

示例实现可能如下所示。首先，我们可以构建一个二进制文件，如下所示(包含随机数据作为占位符)

import numpy as np
from tqdm import tqdm

filename = 'data.bin'
num_samples = 3600000
rows, cols = 30, 32
dtype = np.float32

# format: <num_samples> <rows> <cols> <sample0> <sample1>...
with open(filename, 'wb') as fout:
    # write a header that contains the total number of samples and the rows and columns per sample
    fout.write(np.array((num_samples, rows, cols), dtype=np.int32).tobytes())
    for i in tqdm(range(num_samples)):
        # random placeholder
        sample = np.random.randn(rows, cols).astype(dtype)
        # write data to file
        fout.write(sample.tobytes())

然后我们可以按如下方式定义IterableDataset

import numpy as np
from torch.utils.data import IterableDataset, DataLoader
from tqdm import tqdm

def binary_reader(filename, start=None, end=None, dtype=np.float32):
    itemsize = np.dtype(dtype).itemsize
    with open(filename, 'rb') as fin:
        num_samples, rows, cols = np.frombuffer(fin.read(3 * np.dtype(np.int32).itemsize), dtype=np.int32)
        start = start if start is not None else 0
        end = end if end is not None else num_samples
        blocksize = itemsize * rows * cols
        start_offset = start * blocksize
        fin.seek(start_offset, 1)
        for _ in range(start, end):
            yield np.frombuffer(fin.read(blocksize), dtype=dtype).reshape(rows, cols).copy()


class BinaryIterableDataset(IterableDataset):
    def __init__(self, filename, start=None, end=None, dtype=np.float32):
        super().__init__()
        self.filename = filename
        self.start = start
        self.end = end
        self.dtype = dtype

    def __iter__(self):
        return binary_reader(self.filename, self.start, self.end, self.dtype)

通过在我的系统(使用SSD存储)上快速测试此数据集，我发现我能够在大约10秒内迭代所有360万个样本

dataset = BinaryIterableDataset('data.bin')
for sample in tqdm(dataset):
    pass

3600000it [00:09, 374026.17it/s]

使用DataLoader和batch_size=256大约需要20秒来迭代整个数据集(转换为张量和创建批处理有一些开销)。对于这个数据集，我发现在使用并行加载时向共享内存传输数据和从共享内存传输数据的开销实际上比只使用0个工作器要慢得多。因此，我推荐使用num_workers=0。与任何可迭代的数据集一样，您需要添加额外的逻辑来支持num_worker>；1，尽管我不确定在这种情况下是否值得这样做。

loader = DataLoader(dataset, batch_size=256, num_workers=0)
for batch in tqdm(loader):
    # batch is a tensor of shape (256, 30, 32)
    pass

14063it [00:19, 710.49it/s]

请注意，data.bin文件不能跨使用不同字节顺序的系统移植。但可以进行修改以支持该功能。