Welcome to the TensorDict Documentation!¶
TensorDict is a dictionary-like class that inherits properties from tensors, such as indexing, shape operations, casting to device etc.
You can install tensordict directly from PyPI (see more about installation instructions in the dedicated section below):
$ pip install tensordict
The main purpose of TensorDict is to make code-bases more readable and modular by abstracting away tailored operations:
>>> for i, tensordict in enumerate(dataset):
... # the model reads and writes tensordicts
... tensordict = model(tensordict)
... loss = loss_module(tensordict)
... loss.backward()
... optimizer.step()
... optimizer.zero_grad()
With this level of abstraction, one can recycle a training loop for highly heterogeneous task. Each individual step of the training loop (data collection and transform, model prediction, loss computation etc.) can be tailored to the use case at hand without impacting the others. For instance, the above example can be easily used across classification and segmentation tasks, among many others.
Installation¶
Tensordict releases are synced with PyTorch, so make sure you always enjoy the latest features of the library with the most recent version of PyTorch (although core features are guaranteed to be backward compatible with pytorch>=1.13). Nightly releases can be installed via
$ pip install tensordict-nightly
or via a git clone if you’re willing to contribute to the library:
$ cd path/to/root
$ git clone https://github.com/pytorch/tensordict
$ cd tensordict
$ python setup.py develop