存储和读取海量 3D 数据集的节省空间的方法？ [关闭]

Question

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

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

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

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

最终，我将使用数据文件作为 PyTorch IterableDataset 的输入，因此它必须是可以一次加载一行的内容（例如 .txt 文件，但我希望有一些更好的方法来保存更符合其表格、3 维性质的数据）。任何想法表示赞赏！

Answer 1

由于您计划使用可迭代数据集，因此您不需要随机访问（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 个 worker 慢很多。因此我推荐使用num_workers=0。与任何可迭代数据集一样，您需要添加额外的逻辑来支持 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 对象存储其数据。这种机制允许我们使用FloatStorage.from_file 定义一个内存映射存储系统。使用内存映射张量允许我们将数据集写入磁盘并读取它，就好像它是形状 (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)

这样做的好处是可以兼容内置的TensorDataset

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)