TensorDictBase.map#

TensorDictBase.map(fn: Callable[[TensorCollection], TensorCollection | None], dim: int = 0, num_workers: int | None = None, *, out: TensorCollection | None = None, chunksize: int | None = None, num_chunks: int | None = None, pool: mp.Pool | None = None, generator: Generator | None = None, max_tasks_per_child: int | None = None, worker_threads: int = 1, index_with_generator: bool = False, pbar: bool = False, mp_start_method: str | None = None) → Self#

Applies a function to chunks of the tensordict in worker processes, and gathers the results.

The tensordict is split along dim into chunks, fn is called on each chunk in a pool of worker processes, and the results are joined along dim: concatenated with torch.cat(), or stacked when chunksize=0. For an iterator over the results, see map_iter().

The chunks are slices of the tensordict, not copies. On a memory-mapped tensordict (see memmap_()), the workers read them from the files without a copy, which makes it possible to process data that is larger than memory.

Parameters:

fn (Callable) – the function called on each chunk, a tensordict that holds a slice of the tensordict along dim. It returns a tensordict, a tensor, or None. It must be picklable: define it at the top level of a module or script, not as a lambda.
dim (int, optional) – the dimension along which the tensordict is split into chunks. Defaults to 0.
num_workers (int, optional) – the number of worker processes. Ignored when pool is given. Defaults to None, which uses one worker per CPU.

Keyword Arguments:

out (TensorDictBase, optional) – a tensordict to write the results into, with the batch size of the result. If it is in shared memory or memory-mapped (is_shared() or is_memmap() is True), the workers write into it directly; otherwise the results are sent back and written into it in the main process. Defaults to None.
chunksize (int, optional) – the size of each chunk along dim. 0 passes the elements one by one, unbound along dim. Exclusive with num_chunks. Defaults to None.
num_chunks (int, optional) – the number of chunks. Exclusive with chunksize. Defaults to None: if neither is given, the tensordict is split into one chunk per worker. Use more chunks when a chunk per worker does not fit in memory.
pool (multiprocessing.Pool, optional) – a pool of processes to run the job. When it is given, num_workers, generator, max_tasks_per_child, worker_threads and mp_start_method are ignored. Defaults to None, which creates a pool for the call and closes it afterwards.
generator (Generator, optional) – the generator that draws the base seed of the workers. Worker i seeds torch with the base seed plus i, and NumPy with a seed derived from both, so that the workers draw different random numbers. This makes the workers independent, not the results reproducible: which worker gets which chunk changes from call to call. Defaults to None, which uses a new torch.Generator, so the base seed is the same on every call, and emits a FutureWarning. Starting with TensorDict 0.17, the default will change to drawing the base seed from the global torch RNG, as DataLoader does, which also advances it by one draw. Pass generator=torch.default_generator for that behaviour now, or generator=torch.Generator() to keep the current seeds. To leave the workers unseeded, pass a pool.
max_tasks_per_child (int, optional) – the number of chunks a worker process handles before it is replaced by a new one. Defaults to None, which keeps the workers for the whole call.
worker_threads (int, optional) – the number of torch threads in each worker. Defaults to 1.
index_with_generator (bool, optional) – if True, each chunk is indexed when a worker asks for it, instead of all at once before the job starts. This starts faster, but indexing is slower than split() and chunk(), so the whole call may take longer. Defaults to False.
pbar (bool, optional) – if True, shows a progress bar. Needs tqdm; without it, no bar is shown. Defaults to False.
mp_start_method (str, optional) – the start method of the worker processes, for example "fork" or "spawn". CUDA tensors cannot be shared with processes started with "fork". Defaults to None, which uses the default start method of torch.multiprocessing.

Returns:

The results joined along dim, a tensordict or a tensor depending on what fn returns. When out is given, out. Otherwise, None if fn returns None for every chunk.

Examples

Run this example as a script, not in a REPL: the worker processes must import process_data, and the if __name__ == "__main__" block keeps them from running the main code again.

import torch
from tensordict import TensorDict

def process_data(data):
    # The chunks of a memory-mapped tensordict are locked:
    # return a new tensordict instead of adding a key to data.
    return data.empty().set("y", data["x"] + 1)

if __name__ == "__main__":
    data = TensorDict({"x": torch.zeros(1, 1_000)}, [1, 1_000]).memmap_()
    # generator=torch.default_generator seeds the workers from
    # the global RNG, the default from TensorDict 0.17.
    data = data.map(
        process_data, dim=1, num_workers=2, generator=torch.default_generator
    )
    print(data["y"][:, :10])
    # tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

TensorDictBase.map#

Docs

Tutorials

Resources