Using a dataset in a PyTorch DataLoader#

A deltacells dataset is read through DeltaCellsDataset, which returns batches of any cells you ask for as CSR arrays. This page shows how to put it behind a PyTorch DataLoader. For a complete training setup with PyTorch Lightning, see cellarium-ml, which provides the data loading, models and training loops built for this kind of data.

Reading cells#

import numpy as np
from deltacells import open_dataset

ds = open_dataset("my_dataset")  # a directory, or "gs://bucket/prefix"
ds.n_cells, ds.n_genes, ds.n_tiles

batch = ds.get_batch(np.array([5, 99_000, 12]))  # any global cell indices, in any order
batch.indptr, batch.indices, batch.values  # CSR arrays: int64, int32, float32
batch.to_scipy()  # a scipy.sparse.csr_matrix
batch.to_torch_csr()  # a torch.sparse_csr tensor sharing the memory

Columns are in the stored gene order (genes sorted by decreasing total counts). ds.var_names gives the gene names in that order, and ds.gene_order the original column of each stored gene.

What a loader needs#

a data loader needs	`DeltaCellsDataset` provides
number of cells and genes	`n_cells`, `n_genes`
shard boundaries for shuffling and splitting	`limits`, `tile_bounds(i)`, `tile_of(indices)`
a batch by global cell index	`get_batch(indices)`
the size of a batch, to allocate output buffers (for example in shared memory)	`batch_nnz(indices)`, then `get_batch(indices, out_indices=..., out_values=...)`
“these tiles will be needed next”	`prefetch(tile_ids)`, `prefetch_cells(indices)`
gene names and the gene table	`var_names`, `var`
cell metadata for a batch	`obs.take(indices, columns)`

A minimal DataLoader#

The dataset stores cells in tiles of about 10,000 consecutive cells, and reading a tile is one object read. The efficient way to shuffle is therefore in two levels: visit the tiles in random order, and shuffle the cells within each tile. Batches then come from one tile at a time, and the next tiles can be fetched in the background while the current one is used.

import numpy as np
import torch
from torch.utils.data import DataLoader, IterableDataset, get_worker_info

from deltacells import open_dataset


class TileBatches(IterableDataset):
    """Shuffled batches of cells: tiles in random order, cells shuffled within each tile."""

    def __init__(self, path, batch_size=5000, seed=0):
        self.ds = open_dataset(path)
        self.batch_size, self.seed = batch_size, seed

    def __iter__(self):
        info = get_worker_info()
        worker, n_workers = (info.id, info.num_workers) if info else (0, 1)
        rng = np.random.default_rng([self.seed, worker])
        tiles = np.random.default_rng(self.seed).permutation(self.ds.n_tiles)[worker::n_workers]  # this worker's tiles
        self.ds.prefetch(tiles[:3])
        for k, tile in enumerate(tiles):
            self.ds.prefetch(tiles[k + 1 : k + 3])  # fetch the next tiles while this one is used
            start, stop = self.ds.tile_bounds(int(tile))
            cells = rng.permutation(np.arange(start, stop))
            for i in range(0, len(cells), self.batch_size):
                batch = self.ds.get_batch(cells[i : i + self.batch_size])
                yield {  # plain tensors cross the process boundary; the sparse tensor is built in the training process
                    "indptr": torch.from_numpy(batch.indptr),
                    "indices": torch.from_numpy(batch.indices),
                    "values": torch.from_numpy(batch.values),
                }


def to_csr(batch, n_genes):
    return torch.sparse_csr_tensor(
        batch["indptr"], batch["indices"], batch["values"], size=(len(batch["indptr"]) - 1, n_genes)
    )


# batch_size=None: the dataset already yields whole batches
loader = DataLoader(TileBatches("my_dataset"), batch_size=None, num_workers=4)
n_genes = loader.dataset.ds.n_genes
for batch in loader:
    x = to_csr(batch, n_genes).to("cuda")  # move the sparse batch to the GPU first,
    x = x.to_dense()  # then make it dense there
    ...  # forward and backward pass

Each worker takes its own share of the tiles, so the workers decode in parallel on separate cores. Move the sparse batch to the GPU before densifying it: a dense batch of 5,000 cells by 30,000 genes is 600 MB, while the sparse form of a batch with 2,300 nonzeros per cell is about 90 MB, so the transfer is several times smaller.

Getting the most out of it#

Iterate tile by tile. A batch that spans two tiles reads and decodes both. Shuffling tiles and then cells within a tile (as above) keeps consecutive batches on the same tile, and it is what gives the throughput measured in the benchmarks.
Batches that span many tiles. A batch that touches no more than max_prefetch_tiles tiles (default 8) fetches them all at once. One that touches more, such as a random sample from the whole dataset, is read as a stream: the start of each tile (its indptr, about 8 bytes per cell, one small ranged read) is read first to lay out the output, then the tiles are fetched through a window of max_prefetch_tiles, decoded and gathered one at a time, so memory does not grow with the number of tiles and no tile is read in full twice. This needs a backend with ranged reads, which the local and gs:// backends have (a custom Backend sets ranged_reads = True once its read_range is a real ranged read; otherwise such a batch raises NotImplementedError). It is still slow to read a few cells from each of many tiles, as every tile is read and decoded whole: read the cells of a tile together instead.
Prefetch. prefetch starts fetching and parsing tiles on background threads, which hides the latency of a bucket. The example above looks two tiles ahead.
Memory. A decoded tile takes about 8 bytes per nonzero, which is around 0.4 GB for 10,000 deep cells. max_cached_tiles (default 2) bounds the cache, and decode buffers are reused rather than reallocated.
Threads. decode_threads=1 (the default) is right when each worker has about one core. The decoder releases the Python GIL, so more threads decode a tile’s chunks in parallel.
Pickling and forking. The dataset object can be pickled, which is how it reaches the DataLoader workers when they are started by spawn. Caches and threads are rebuilt in each worker. Workers started by fork (the default on Linux) get the same fresh state: a forked process starts with empty caches and no threads, so it is safe to read from the dataset in the main process before creating the DataLoader. A GCSBackend likewise makes its own client in a forked process.
Fewer copies. To avoid one copy between the worker and the training process, allocate the output tensors in shared memory and pass them to get_batch (batch_nnz gives the size). The reference loader does this.

Cell metadata#

If the dataset has cell metadata, columns are read separately from the counts, so a loader that does not need them pays nothing. Choose the columns up front and they are copied once into a local cache:

ds = open_dataset("my_dataset", obs_columns=["cell_type", "donor_id"])
codes = ds.obs.take(batch_cell_indices, ["cell_type", "donor_id"])  # integer codes for the batch, a local lookup
ds.obs.categories("cell_type")  # the vocabulary behind the codes

See Cell metadata (obs) and genes (var) for how the metadata is stored, cached and read.

A fuller example#

The repository has two examples, both runnable:

examples/reference_loader.py is a complete IterableDataset with tile shuffling per epoch, splitting across data-parallel replicas and workers, prefetch hints and shared-memory output.
examples/dataloader_example.py runs it end to end on synthetic data.

These show the pattern; they are examples, not part of the package. For training, use cellarium-ml.

Using a dataset in a PyTorch DataLoader

Contents