Space efficient way to store and read massive 3d dataset?
I am trying to train a neural network on sequential data. My dataset will consist of 3.6 million training examples. Each example will be a 30 x 32 ndarray (32 features observed over 30 days).
My question is what is the most space-efficient way to write and read this data?
Essentially it will have shape (3.6m, 30, 32) and np.save() seems convenient but I can't hold this whole thing in memory so I can't really save it using np.save() (or load it back using np.load()). CSV also won't work because my data has 3 dimensions.
My plan to create the thing is to process entries in batches and append them to some file so that I can keep memory free as I go.
Eventually, I am going to use the data file as an input for a PyTorch IterableDataset so it must be something that can be loaded one line at a time (like a .txt file, but I'm hoping there is some better way to save this data that is more true to its tabular, 3-dimensional nature). Any ideas are appreciated!

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: &lt;num_samples&gt; &lt;rows&gt; &lt;cols&gt; &lt;sample0&gt; &lt;sample1&gt;...
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 &gt; 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.
