存储和读取海量3d数据集的节省空间的方法？

Question

我正试图在序列数据上训练神经网络。我的数据集将包含360万个培训示例。每个示例将是一个30 x 32 Ndaray（在30天内观察32个特征）

我的问题是写和读这些数据最节省空间的方式是什么

从本质上讲，它将具有(3.6m, 30, 32)的形状，而且np.save()看起来很方便，但我无法将整个内容保存在内存中，因此我无法使用np.save()保存它（或者使用np.load()将其加载回去）。CSV也不起作用，因为我的数据有3个维度

我创建这个东西的计划是成批处理条目并将它们附加到某个文件中，这样我就可以在运行时保持内存空闲

最后，我将使用数据文件作为PyTorch IterableDataset的输入，因此它必须是可以一次加载一行的文件（比如.txt文件，但我希望有更好的方法来保存这些数据，使其更符合表格、三维的性质）。任何想法都很感激

Answer 1

由于您计划使用iterable数据集，因此不需要随机访问（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]

使用带有batch_size=256的DataLoader迭代整个数据集大约需要20秒（转换为张量和创建批处理会有一些开销）。对于这个数据集，我发现当使用并行加载时，将数据传输到共享内存和从共享内存传输数据的开销实际上比仅使用0个worker要慢得多。因此，我建议使用num_workers=0。与任何iterable数据集一样，您需要添加额外的逻辑来支持num_workers>；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文件不能跨使用不同字节顺序的系统移植。尽管可以进行修改以支持这一点

Answer 2

另一种解决方案是使用内存映射张量。这与other solution类似，但在IMO中更好，因为它抽象掉了与二进制数据的直接交互，并在更高的抽象级别上运行

每个张量使用Storage对象存储其数据。此机制允许我们使用^{}定义内存映射存储系统。使用内存映射的张量，我们既可以将数据集写入磁盘，也可以像正常的形状张量（3600000、32、30）一样读取数据集，而无需将内存直接存储在RAM中

例如，我们可以使用如下方法将数据集写入磁盘

import torch

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

# shared=True allows us to save the tensor to disk as we perform in place modifications to it
samples = torch.FloatTensor(torch.FloatStorage.from_file(filename, shared=True, size=num_samples * rows * cols)).reshape(num_samples, rows, cols)

for idx in tqdm(range(num_samples)):
    # placeholder random samples, insert your actual samples here
    # every in-place assignment to samples is automatically reflected on the disk
    samples[idx] = torch.randn(rows, cols)

这样做的好处是与内置的^{}兼容

from torch.utils.data import TensorDataset, DataLoader

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

# shared=False prevents changes to samples from affecting the data on disk
samples = torch.FloatTensor(torch.FloatStorage.from_file(filename, shared=False, size=num_samples * rows * cols)).reshape(num_samples, rows, cols)

dataset = TensorDataset(samples)
loader = DataLoader(dataset, batch_size=256, num_workers=0)

for batch in tqdm(loader):
    # batch is a (256, 32, 30) tensor
    pass

100%|██████████| 14063/14063 [00:11<00:00, 1216.80it/s]

Answer 3

而不是np.save：

np.save('data.npy', x)
retrieved_array = np.load('data.npy')

您可以使用：

np.savez_compressed('data.npz', array=x)
retrieved_array = np.load('data.npz')['array']

这有助于将我的笔记本电脑上的数据大小从375MB减少到60MB，具体数据如下：

x = np.random.randint(0,10, size=(30, 32, 100000))

备注1：请注意，这不是高效的：

x = np.random.randint(0,10, size=(30,32,10000))
%timeit np.save('data.npy', x)
%timeit np.load('data.npy')
%timeit np.savez_compressed('data.npz', array=x)
%timeit np.load('data.npz')['array']

153 ms ± 42.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
35.7 ms ± 3.59 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
2.57 s ± 209 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
163 ms ± 18.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)