Stochastic gradient descent wants every minibatch to look like a random sample of the whole dataset. When the data fits in memory this is free. When it sits on disk, every randomly chosen sample is a separate small read, and disks are slow at small random reads. Single-cell atlases are the case that made me write scDataset: Tahoe-100M has 100 million cells in 14 AnnData files.
scDataset combines two ideas. Block sampling shuffles blocks of b contiguous samples instead of single samples, so a minibatch of m samples needs m/b reads instead of m. Batched fetching reads f minibatches’ worth of blocks in one call, shuffles them in memory and splits them into f minibatches, so the samples of a block are spread across several minibatches.
On Tahoe-100M this reads data more than a hundred times faster than picking cells at random, and models trained on it learn as well as with true random sampling. Reading the files in order does much worse, even through a shuffle buffer, because each long stretch of a file comes from a single experiment.
It reads AnnData files directly, without converting them, and works with any data that can be indexed by a list of positions: NumPy arrays, HuggingFace datasets, PyTorch datasets.
from scdataset import scDataset, BlockShuffling
from torch.utils.data import DataLoader
strategy = BlockShuffling(block_size=16)
dataset = scDataset(data, strategy, batch_size=64, fetch_factor=256)
loader = DataLoader(dataset, batch_size=None, num_workers=4)Install it with pip install scDataset. The documentation covers the other sampling strategies, multiple files and the transform hooks.