Datasets that do not fit in memory

Setup#

pip install torch numpy

Written and tested against torch 2.5 and NumPy 1.26 on Windows. The file created is a modest 32 MB — big enough to be honest, small enough to be polite.

A memory-mapped Dataset#

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

# One-time step: write 200,000 rows to disk (~30 MB). Real projects do this once.
path = "big_features.npy"
rows, cols = 200_000, 40
np.save(path, np.random.default_rng(0).standard_normal((rows, cols), dtype=np.float32))

class MemmapDataset(Dataset):
    """Reads rows straight from disk. RAM holds one batch, not the file."""
    def __init__(self, path):
        self.data = np.load(path, mmap_mode="r")     # opens the file, copies nothing

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        row = np.array(self.data[idx])               # this line reads ONE row from disk
        return torch.from_numpy(row)

ds = MemmapDataset(path)
print("rows on disk:", len(ds))
print("one sample:", tuple(ds[123].shape), ds[123].dtype)

loader = DataLoader(ds, batch_size=256, shuffle=True,
                    generator=torch.Generator().manual_seed(0))
x = next(iter(loader))
print("one batch:", tuple(x.shape))
import os
print(f"file size: {os.path.getsize(path) / 1e6:.0f} MB")

rows on disk: 200000
one sample: (40,) torch.float32
one batch: (256, 40)
file size: 32 MB

At 32 MB this file fits in RAM easily — the point is that nothing in the code would change at 320 GB. That is the property you are buying.

The walkthrough#

mmap_mode="r" is the entire magic. np.load without it reads the whole file into RAM; with it, the call returns instantly having read only a small header, and the array is a window onto the file. Indexing a row makes the operating system fetch the pages containing that row — a few kilobytes — and nothing else.

np.array(self.data[idx]) forces a real copy of the one row. Skip it and you return a live view into the map; a batch of views keeps the map pinned across process boundaries and interacts messily with workers. Copy the row; it is tiny.

shuffle=True works untouched — this is the quiet advantage over streaming. The map preserves random access, so the shelf semantics — shuffling, samplers, exact epochs — all survive. Streams give none of that back.

Building the file in the first place: convert once, in chunks — read 10,000 rows from the source (CSVs, a database), write them into a pre-allocated np.lib.format.open_memmap(...) array, repeat. The conversion script also fits in constant RAM.

Storage formats, honestly ranked for this job#

Raw .npy + memmap: unbeatable for fixed-size numeric rows — zero decode cost, perfect random access. Fixed size is the catch: variable-length samples need either padding on disk or an offsets file alongside. Images usually stay as individual encoded files (JPEG is its own compression; a folder is an out-of-core store, as the lazy-loading pattern showed). Parquet/Arrow suit tabular work with column selection. HDF5 handles groups of arrays but carries a sharp multiprocessing edge, noted below.

Common mistakes#

Benchmarking on warm cache. The second epoch reads from the operating system's page cache — RAM — and flies. The first epoch off a cold disk tells the truth. Judge speed on epoch one, or after a reboot.

Random access on a spinning hard drive. Memmap plus shuffle=True on an HDD means a physical head seek per row — brutal. On SSDs this pattern is fine; on HDDs, coarsen randomness: shuffle chunks of a few thousand contiguous rows, then shuffle within the loaded chunk.

HDF5 handles crossing into workers. An h5py.File opened in __init__ breaks when workers fork or spawn — corrupt reads or crashes. The fix is lazy per-worker opening: store the path, open on first __getitem__ call in each worker (if self.f is None: self.f = h5py.File(...)).

dtype surprises doubling the file. Saving float64 doubles size and halves throughput for no modelling gain. Decide dtype at conversion — float32, or even float16/uint8 where the data allows — and write it once.

Try it yourself#

Time one full pass at batch_size=256 twice in the same run — cold-ish, then warm. The gap you measure is the page cache at work, and it is exactly why the "benchmark warm" mistake fools people.

What to learn next#

• Surviving corrupt and missing files — what happens when row 4,081,337 turns out to be damaged.
• IterableDataset for streams and huge files — the other answer, when even random access must go.
• Finding out whether the GPU is waiting for data — checking whether your disk reads keep up.