tensordict package

The TensorDict class simplifies the process of passing multiple tensors from module to module by packing them in a dictionary-like object that inherits features from regular pytorch tensors.

TensorDictBase()	TensorDictBase is an abstract parent class for TensorDicts, a torch.Tensor data container.
TensorDict([source, batch_size, device, ...])	A batched dictionary of tensors.
LazyStackedTensorDict(*tensordicts[, ...])	A Lazy stack of TensorDicts.
PersistentTensorDict(*[, batch_size, ...])	Persistent TensorDict implementation.
TensorDictParams(parameters, *[, ...])	Holds a TensorDictBase instance full of parameters.

TensorDict as a context manager

TensorDict can be used as a context manager in situations where an action has to be done and then undone. This include temporarily locking/unlocking a tensordict

>>> data.lock_()  # data.set will result in an exception
>>> with data.unlock_():
...     data.set("key", value)
>>> assert data.is_locked()

or to execute functional calls with a TensorDict instance containing the parameters and buffers of a model:

>>> params = TensorDict.from_module(module).clone()
>>> params.zero_()
>>> with params.to_module(module):
...     y = module(x)

In the first example, we can modify the tensordict data because we have temporarily unlocked it. In the second example, we populate the module with the parameters and buffers contained in the params tensordict instance, and reset the original parameters after this call is completed.

Memory-mapped tensors

tensordict offers the MemoryMappedTensor primitive which allows you to work with tensors stored in physical memory in a handy way. The main advantages of MemoryMappedTensor are its ease of construction (no need to handle the storage of a tensor), the possibility to work with big contiguous data that would not fit in memory, an efficient (de)serialization across processes and efficient indexing of stored tensors.

If all workers have access to the same storage (both in multiprocess and distributed settings), passing a MemoryMappedTensor will just consist in passing a reference to a file on disk plus a bunch of extra meta-data for reconstructing it. The same goes with indexed memory-mapped tensors as long as the data-pointer of their storage is the same as the original one.

Indexing memory-mapped tensors is much faster than loading several independent files from the disk and does not require to load the full content of the array in memory. However, physical storage of PyTorch tensors should not be any different:

>>> my_images = MemoryMappedTensor.empty((1_000_000, 3, 480, 480), dtype=torch.unint8)
>>> mini_batch = my_images[:10]  # just reads the first 10 images of the dataset

MemoryMappedTensor(source, *[, dtype, ...])

A Memory-mapped Tensor.

Utils

utils.expand_as_right(tensor, dest)	Expand a tensor on the right to match another tensor shape.
utils.expand_right(tensor, shape)	Expand a tensor on the right to match a desired shape.
utils.isin(input, reference, key[, dim])	Tests if each element of key in input dim is also present in the reference.
utils.remove_duplicates(input, key[, dim, ...])	Removes indices duplicated in key along the specified dimension.
is_batchedtensor(arg0)
is_tensor_collection(datatype)	Checks if a data object or a type is a tensor container from the tensordict lib.
make_tensordict([input_dict, batch_size, device])	Returns a TensorDict created from the keyword arguments or an input dictionary.
merge_tensordicts(*tensordicts)	Merges tensordicts together.
pad(tensordict, pad_size[, value])	Pads all tensors in a tensordict along the batch dimensions with a constant value, returning a new tensordict.
pad_sequence(list_of_tensordicts[, ...])	Pads a list of tensordicts in order for them to be stacked together in a contiguous format.
dense_stack_tds(td_list[, dim])	Densely stack a list of TensorDictBase objects (or a LazyStackedTensorDict) given that they have the same structure.
set_lazy_legacy(mode)	Sets the behaviour of some methods to a lazy transform.
lazy_legacy([allow_none])	Returns True if lazy representations will be used for selected methods.