Space efficient way to store and read massive 3d dataset?

An alternative solution is to use memory mapped tensors. This is similar to the other solution but is better IMO since it abstracts away the direct interaction with binary data and operates at a higher level of abstraction.

Every tensor stores its data using a Storage object. This mechanism allows us to define a memory mapped storage system using FloatStorage.from_file. Using a memory mapped tensor allows us to both write our dataset to disk and read it as if it were a normal tensor of shape (3600000, 32, 30) without directly storing that memory in RAM.

For example we could write our dataset to disk using something like the following

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)

This has the benefit of being compatible with the built-in 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]

Since you are planning on using an iterable dataset you shouldn't need random access (IterableDataset doesn't support shuffle samplers). In that case, why not just write everything to a binary file and iterate over that? I find in practice this often is much faster than alternative solutions. This should be much faster than saving as a text file since you avoid the overhead of converting text to numbers.

An example implementation may look something like the following. First we could build a binary file as follows (containing random data as a placeholder)

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())

Then we could define an IterableDataset as follows

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)

From a quick test of this dataset on my system (which uses SSD storage) I find I am able to iterate over all 3.6 million samples in about 10 seconds

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

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

Using a DataLoader with batch_size=256 it takes me about 20 seconds to iterate over the whole dataset (converting to tensors and creating batches has some overhead). For this dataset I found that the overhead of transferring data to and from shared memory when using parallel loading is actually quite a bit slower than just using 0 workers. Therefore I recommend using num_workers=0. As with any iterable dataset you would need to add extra logic to support num_workers > 1, though I'm not sure it would be worth it in this case.

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]

Note that the data.bin file would not be portable across systems that use different byte order. Though modifications could be made to support that.

Instead of np.save:

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

You could use:

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

It helped to reduce size of data from 375MB to 60MB on my laptop on this data:

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

Remark1: note that this is not time efficient:

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)