Load data and share it between processes#

A tensordict with a batch dimension is a dataset. A DataLoader reads each batch from it with one index operation, and the batch is a tensordict. Processes on one node can also share one tensordict, in shared memory or in memory-mapped files, and see the writes of the other processes. For data from Hugging Face datasets and for training with Lightning, see Use tensordict with other libraries.

Use a tensordict as a dataset#

A tensordict is a map-style dataset: len(td) is the size of its first batch dimension, and td[i] is one sample (see Batch size, shapes, and indexing). It also has __getitems__, which reads a whole batch with one index operation. Pass collate_fn=lambda x: x to keep that batch as it is:

>>> import torch
>>> from torch.utils.data import DataLoader
>>> from tensordict import TensorDict
>>> _ = torch.manual_seed(0)
>>> data = TensorDict(
...     {"image": torch.randn(10, 1, 8, 8), "label": torch.arange(10)}, batch_size=[10]
... )
>>> loader = DataLoader(data, batch_size=4, shuffle=True, collate_fn=lambda x: x)
>>> batches = list(loader)
>>> [batch["label"].tolist() for batch in batches]
[[5, 1, 2, 7], [9, 4, 6, 3], [0, 8]]
>>> type(batches[0]).__name__, batches[0].batch_size
('TensorDict', torch.Size([4]))

The default collate_fn returns that batch as it is too:

>>> next(iter(DataLoader(data, batch_size=4))).batch_size
torch.Size([4])

The same works for a memory-mapped tensordict (see Save and load tensordicts). Each batch then reads only its rows from disk.

Without a DataLoader, index the tensordict with a permutation. With no worker processes, this gives the same batches:

>>> for idx in torch.randperm(len(data)).split(4):
...     batch = data[idx]
>>> batch.batch_size
torch.Size([2])

A dataset that returns tensordicts#

A dataset can return one tensordict for each sample, for example a list of tensordicts, or a Dataset whose __getitem__ builds a tensordict. The DataLoader then gets a list of tensordicts for each batch. The default collate_fn stacks them into one tensordict, as collate() does (see below). A batch of tensorclasses keeps its class:

>>> samples = [TensorDict({"features": torch.randn(3), "label": i}) for i in range(10)]
>>> batch = next(iter(DataLoader(samples, batch_size=4)))
>>> batch.batch_size, batch["label"]
(torch.Size([4]), tensor([0, 1, 2, 3]))
>>> from tensordict import TensorClass
>>> class Sample(TensorClass):
...     features: torch.Tensor
>>> batch = next(iter(DataLoader([Sample(features=torch.randn(3)) for _ in range(10)], batch_size=4)))
>>> type(batch).__name__, batch.features.shape
('Sample', torch.Size([4, 3]))

A batch that mixes different tensorclasses raises a TypeError, unless one class is a base of the others and has the same fields. The batch then takes the class of its first sample.

collate_fn=torch.stack gives the same batches:

>>> batch = next(iter(DataLoader(samples, batch_size=4, collate_fn=torch.stack)))
>>> batch.batch_size, batch["label"]
(torch.Size([4]), tensor([0, 1, 2, 3]))

If the samples have different shapes, torch.stack raises an error. Pass lazy_stack() instead. It returns a LazyStackedTensorDict, which keeps each sample as it is (see Stacking and variable-size data). This example loads point clouds with different numbers of points:

>>> from tensordict import lazy_stack
>>> clouds = [TensorDict({"points": torch.randn(n, 3)}) for n in (5, 2, 7, 3)]
>>> batch = next(iter(DataLoader(clouds, batch_size=4, collate_fn=lazy_stack)))
>>> type(batch).__name__, batch.batch_size
('LazyStackedTensorDict', torch.Size([4]))
>>> [cloud["points"].shape for cloud in batch.unbind(0)]
[torch.Size([5, 3]), torch.Size([2, 3]), torch.Size([7, 3]), torch.Size([3, 3])]

collate() does both: it stacks the samples densely if their shapes match, and lazily otherwise. It also reads samples that hold tensordicts or tensorclasses in a tuple or a dict, such as (td, label) pairs, which torch.stack cannot stack:

>>> from tensordict import collate
>>> pairs = [(TensorDict({"features": torch.randn(3)}), i) for i in range(10)]
>>> features, labels = next(iter(DataLoader(pairs, batch_size=4, collate_fn=collate)))
>>> type(features).__name__, features.batch_size, labels
('TensorDict', torch.Size([4]), tensor([0, 1, 2, 3]))
>>> batch = next(iter(DataLoader(clouds, batch_size=4, collate_fn=collate)))
>>> type(batch).__name__
'LazyStackedTensorDict'

The default collate_fn gives the same batches, so you need to pass collate() only for its keyword arguments, shown below:

>>> features, labels = next(iter(DataLoader(pairs, batch_size=4)))
>>> type(features).__name__, labels
('TensorDict', tensor([0, 1, 2, 3]))
>>> type(next(iter(DataLoader(clouds, batch_size=4)))).__name__
'LazyStackedTensorDict'

To pad some entries instead, each with its own value, or to stack them as jagged nested tensors, pass the keyword arguments of collate() with functools.partial(). The other entries are collated as above:

>>> from functools import partial
>>> tokens = [TensorDict({"ids": torch.ones(n, dtype=torch.long), "points": torch.randn(n, 3)}) for n in (2, 4, 3)]
>>> collate_fn = partial(collate, pad_keys={"ids": 7}, nested_keys=["points"], mask_key="mask")
>>> batch = next(iter(DataLoader(tokens, batch_size=3, collate_fn=collate_fn)))
>>> batch["ids"].shape, batch["mask", "ids"].sum(-1), batch["points"].is_nested
(torch.Size([3, 4]), tensor([2, 4, 3]), True)

Batch samplers#

A batch sampler gives the indices of a whole batch. Pass it as batch_sampler. The DataLoader then calls data.__getitems__ with the list of indices of each batch, and collate_fn gets the tensordict:

>>> from torch.utils.data import BatchSampler, SequentialSampler
>>> sampler = BatchSampler(SequentialSampler(data), batch_size=4, drop_last=True)
>>> loader = DataLoader(data, batch_sampler=sampler, collate_fn=lambda x: x)
>>> [batch["label"].tolist() for batch in loader]
[[0, 1, 2, 3], [4, 5, 6, 7]]

Passing it as sampler, with batch_size=None, gives the same batches through data[indices], but the workers of the DataLoader then send an HDF5 or zarr PersistentTensorDict batch as a lazy view, which the main process reads.

Worker processes and pinned memory#

With num_workers greater than 0, the DataLoader reads the batches in worker processes. Each worker gets the dataset when it starts:

With the fork start method (the default on Linux up to Python 3.13), the worker inherits a copy-on-write view of the memory of the main process.
With spawn (the default on macOS and Windows) and forkserver (the default on Linux from Python 3.14), the DataLoader pickles the dataset to each worker. A memory-mapped tensordict pickles as its file names, so the workers open the same files and do not copy the data. Pickle cannot send a lambda, so define collate_fn as a function at the top level of a module.

In a worker, __getitems__ of an HDF5 or zarr PersistentTensorDict reads the rows into a TensorDict in memory, so the batch is not a view that writes to the file, as it is without workers. With persistent_workers=True, each worker keeps its own HDF5 file handle across epochs, so it may not see writes that the main process makes to the file between epochs: HDF5 does not support writes across processes without SWMR mode.

With spawn or forkserver, each worker opens the HDF5 file again. While the main process has the file open for writing (mode="r+", "a" or "w"), HDF5’s file lock stops the workers from opening it, so they fail with BlockingIOError, and the DataLoader raises RuntimeError saying that a worker exited unexpectedly. Give the DataLoader a tensordict opened with mode="r", and don’t keep the file open for writing while the workers run; or use fork. With persistent_workers=True, the workers keep the file open for reading across epochs, so the main process cannot open it for writing between epochs either.

The worker examples start processes, so they are complete scripts. Save one to a file and run it with python:

import tempfile

import torch
from torch.utils.data import DataLoader

from tensordict import TensorDict


def keep_batch(batch):
    return batch


if __name__ == "__main__":
    with tempfile.TemporaryDirectory() as path:
        data = TensorDict(
            {
                "image": torch.randn(1_000, 1, 8, 8),
                "label": torch.arange(1_000),
            },
            batch_size=[1_000],
        ).memmap_(path)
        loader = DataLoader(
            data, batch_size=100, shuffle=True, num_workers=2, collate_fn=keep_batch
        )
        total = sum(batch["label"].sum().item() for batch in loader)
        print(total)  # 499500

With pin_memory=True, the DataLoader calls pin_memory() on each batch, which pins every leaf. From pinned memory, to(device, non_blocking=True) copies the batch to the GPU while the CPU goes on. These lines need a GPU, so they are not run when this page is tested:

>>> loader = DataLoader(data, batch_size=4, collate_fn=lambda x: x, pin_memory=True)  
>>> batch = next(iter(loader)).to("cuda", non_blocking=True)  

Without a DataLoader, td.to("cuda", non_blocking_pin=True) pins each leaf on the calling thread and copies it to the GPU (see Devices, dtypes, and memory). For leaves of a few KiB, a plain td.to("cuda") is usually faster.

Several processes on one dataset on disk#

Processes can also open a memory-mapped dataset by its path, with load_memmap(). Readers that do not change the data can pass mode="r" (see Save and load tensordicts). tensordict does not lock the files, so give each writer its own rows (see Share data between processes). This script preallocates a dataset on disk, fills it from two processes, and reads it back:

import tempfile

import torch
import torch.multiprocessing as mp

from tensordict import TensorDict


def fill(path, start, stop):
    data = TensorDict.load_memmap(path)
    for i in range(start, stop):
        data[i] = TensorDict({"image": torch.full((3, 8, 8), i, dtype=torch.uint8)})


if __name__ == "__main__":
    with tempfile.TemporaryDirectory() as path:
        TensorDict.from_schema(
            {"image": ([3, 8, 8], torch.uint8)},
            batch_size=[100],
            storage="memmap",
            prefix=path,
        )
        writers = [
            mp.Process(target=fill, args=(path, 0, 50)),
            mp.Process(target=fill, args=(path, 50, 100)),
        ]
        for p in writers:
            p.start()
        for p in writers:
            p.join()
        data = TensorDict.load_memmap(path, mode="r")
        # tensor([ 0, 49, 50, 99], dtype=torch.uint8)
        print(data["image"][[0, 49, 50, 99], 0, 0, 0])

Process a dataset in parallel with map#

map() splits a tensordict into chunks along one dimension, calls a function on each chunk in a pool of processes, and concatenates the results. If the tensordict is memory-mapped or in shared memory, the chunks go to the workers without a copy. Return a new tensordict from the function, and do not add keys to its input: the chunk of a memory-mapped tensordict is locked. If you pass a preallocated memory-mapped tensordict as out, the workers write the results into it. This script computes the kinetic energy of 8 particles in two processes:

import tempfile

import torch

from tensordict import TensorDict


def kinetic_energy(chunk):
    energy = 0.5 * chunk["mass"] * (chunk["vel"] ** 2).sum(-1)
    return TensorDict({"energy": energy}, batch_size=chunk.batch_size)


if __name__ == "__main__":
    with tempfile.TemporaryDirectory() as path:
        particles = TensorDict(
            {
                "mass": torch.ones(8),
                "vel": torch.arange(8.0).unsqueeze(-1).repeat(1, 3),
            },
            batch_size=[8],
        ).memmap_()
        out = TensorDict({"energy": torch.zeros(8)}, batch_size=[8]).memmap_like(path)
        particles.map(
            kinetic_energy,
            num_workers=2,
            chunksize=2,
            out=out,
            generator=torch.default_generator,
        )
        # MemoryMappedTensor([ 0.0000,  1.5000,  6.0000, 13.5000])
        print(out["energy"][:4])

map also takes num_chunks, a pool that you create, and a generator that seeds each worker. Starting with TensorDict 0.17, the default will change to drawing the seed from the global RNG; generator=torch.default_generator does that now. map_iter() returns the results one chunk at a time.

Docs

Tutorials

Resources