Skip to main content

No project description provided

Project description

Docs - GitHub.io Benchmarks Python version GitHub license pypi version pypi nightly version Downloads Downloads codecov circleci Conda - Platform Conda (channel only)

TensorDict

Installation | General features | Tensor-like features | Distributed capabilities | TensorDict for functional programming using FuncTorch | Lazy preallocation | Nesting TensorDicts | TensorClass

TensorDict is a dictionary-like class that inherits properties from tensors, such as indexing, shape operations, casting to device or point-to-point communication in distributed settings.

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.

Features

General

A tensordict is primarily defined by its batch_size (or shape) and its key-value pairs:

>>> from tensordict import TensorDict
>>> import torch
>>> tensordict = TensorDict({
...     "key 1": torch.ones(3, 4, 5),
...     "key 2": torch.zeros(3, 4, 5, dtype=torch.bool),
... }, batch_size=[3, 4])

The batch_size and the first dimensions of each of the tensors must be compliant. The tensors can be of any dtype and device. Optionally, one can restrict a tensordict to live on a dedicated device, which will send each tensor that is written there:

>>> tensordict = TensorDict({
...     "key 1": torch.ones(3, 4, 5),
...     "key 2": torch.zeros(3, 4, 5, dtype=torch.bool),
... }, batch_size=[3, 4], device="cuda:0")
>>> tensordict["key 3"] = torch.randn(3, 4, device="cpu")
>>> assert tensordict["key 3"].device is torch.device("cuda:0")

Tensor-like features

TensorDict objects can be indexed exactly like tensors. The resulting of indexing a TensorDict is another TensorDict containing tensors indexed along the required dimension:

>>> tensordict = TensorDict({
...     "key 1": torch.ones(3, 4, 5),
...     "key 2": torch.zeros(3, 4, 5, dtype=torch.bool),
... }, batch_size=[3, 4])
>>> sub_tensordict = tensordict[..., :2]
>>> assert sub_tensordict.shape == torch.Size([3, 2])
>>> assert sub_tensordict["key 1"].shape == torch.Size([3, 2, 5])

Similarly, one can build tensordicts by stacking or concatenating single tensordicts:

>>> tensordicts = [TensorDict({
...     "key 1": torch.ones(3, 4, 5),
...     "key 2": torch.zeros(3, 4, 5, dtype=torch.bool),
... }, batch_size=[3, 4]) for _ in range(2)]
>>> stack_tensordict = torch.stack(tensordicts, 1)
>>> assert stack_tensordict.shape == torch.Size([3, 2, 4])
>>> assert stack_tensordict["key 1"].shape == torch.Size([3, 2, 4, 5])
>>> cat_tensordict = torch.cat(tensordicts, 0)
>>> assert cat_tensordict.shape == torch.Size([6, 4])
>>> assert cat_tensordict["key 1"].shape == torch.Size([6, 4, 5])

TensorDict instances can also be reshaped, viewed, squeezed and unsqueezed:

>>> tensordict = TensorDict({
...     "key 1": torch.ones(3, 4, 5),
...     "key 2": torch.zeros(3, 4, 5, dtype=torch.bool),
... }, batch_size=[3, 4])
>>> print(tensordict.view(-1))
torch.Size([12])
>>> print(tensordict.reshape(-1))
torch.Size([12])
>>> print(tensordict.unsqueeze(-1))
torch.Size([3, 4, 1])

One can also send tensordict from device to device, place them in shared memory, clone them, update them in-place or not, split them, unbind them, expand them etc.

If a functionality is missing, it is easy to call it using apply() or apply_():

tensordict_uniform = tensordict.apply(lambda tensor: tensor.uniform_())

Distributed capabilities

Complex data structures can be cumbersome to synchronize in distributed settings. tensordict solves that problem with synchronous and asynchronous helper methods such as recv, irecv, send and isend that behave like their torch.distributed counterparts:

>>> # on all workers
>>> data = TensorDict({"a": torch.zeros(()), ("b", "c"): torch.ones(())}, [])
>>> # on worker 1
>>> data.isend(dst=0)
>>> # on worker 0
>>> data.irecv(src=1)

When nodes share a common scratch space, the MemmapTensor backend can be used to seamlessly send, receive and read a huge amount of data.

TensorDict for functional programming using FuncTorch

We also provide an API to use TensorDict in conjunction with FuncTorch. For instance, TensorDict makes it easy to concatenate model weights to do model ensembling:

>>> from torch import nn
>>> from tensordict import TensorDict
>>> from tensordict.nn import make_functional
>>> import torch
>>> from torch import vmap
>>> layer1 = nn.Linear(3, 4)
>>> layer2 = nn.Linear(4, 4)
>>> model = nn.Sequential(layer1, layer2)
>>> # we represent the weights hierarchically
>>> weights1 = TensorDict(layer1.state_dict(), []).unflatten_keys(".")
>>> weights2 = TensorDict(layer2.state_dict(), []).unflatten_keys(".")
>>> params = make_functional(model)
>>> assert (params == TensorDict({"0": weights1, "1": weights2}, [])).all()
>>> # Let's use our functional module
>>> x = torch.randn(10, 3)
>>> out = model(x, params=params)  # params is the last arg (or kwarg)
>>> # an ensemble of models: we stack params along the first dimension...
>>> params_stack = torch.stack([params, params], 0)
>>> # ... and use it as an input we'd like to pass through the model
>>> y = vmap(model, (None, 0))(x, params_stack)
>>> print(y.shape)
torch.Size([2, 10, 4])

Moreover, tensordict modules are compatible with torch.fx and torch.compile, which means that you can get the best of both worlds: a codebase that is both readable and future-proof as well as efficient and portable!

Lazy preallocation

Pre-allocating tensors can be cumbersome and hard to scale if the list of preallocated items varies according to the script configuration. TensorDict solves this in an elegant way. Assume you are working with a function foo() -> TensorDict, e.g.

def foo():
    tensordict = TensorDict({}, batch_size=[])
    tensordict["a"] = torch.randn(3)
    tensordict["b"] = TensorDict({"c": torch.zeros(2)}, batch_size=[])
    return tensordict

and you would like to call this function repeatedly. You could do this in two ways. The first would simply be to stack the calls to the function:

tensordict = torch.stack([foo() for _ in range(N)])

However, you could also choose to preallocate the tensordict:

tensordict = TensorDict({}, batch_size=[N])
for i in range(N):
    tensordict[i] = foo()

which also results in a tensordict (when N = 10)

TensorDict(
    fields={
        a: Tensor(torch.Size([10, 3]), dtype=torch.float32),
        b: TensorDict(
            fields={
                c: Tensor(torch.Size([10, 2]), dtype=torch.float32)},
            batch_size=torch.Size([10]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

When i==0, your empty tensordict will automatically be populated with empty tensors of batch-size N. After that, updates will be written in-place. Note that this would also work with a shuffled series of indices (pre-allocation does not require you to go through the tensordict in an ordered fashion).

Nesting TensorDicts

It is possible to nest tensordict. The only requirement is that the sub-tensordict should be indexable under the parent tensordict, i.e. its batch size should match (but could be longer than) the parent batch size.

We can switch easily between hierarchical and flat representations. For instance, the following code will result in a single-level tensordict with keys "key 1" and "key 2.sub-key":

>>> tensordict = TensorDict({
...     "key 1": torch.ones(3, 4, 5),
...     "key 2": TensorDict({"sub-key": torch.randn(3, 4, 5, 6)}, batch_size=[3, 4, 5])
... }, batch_size=[3, 4])
>>> tensordict_flatten = tensordict.flatten_keys(separator=".")

Accessing nested tensordicts can be achieved with a single index:

>>> sub_value = tensordict["key 2", "sub-key"]

TensorClass

Content flexibility comes at the cost of predictability. In some cases, developers may be looking for data structure with a more explicit behavior. tensordict provides a dataclass-like decorator that allows for the creation of custom dataclasses that support the tensordict operations:

>>> from tensordict.prototype import tensorclass
>>> import torch
>>>
>>> @tensorclass
... class MyData:
...    image: torch.Tensor
...    mask: torch.Tensor
...    label: torch.Tensor
...
...    def mask_image(self):
...        return self.image[self.mask.expand_as(self.image)].view(*self.batch_size, -1)
...
...    def select_label(self, label):
...        return self[self.label == label]
...
>>> images = torch.randn(100, 3, 64, 64)
>>> label = torch.randint(10, (100,))
>>> mask = torch.zeros(1, 64, 64, dtype=torch.bool).bernoulli_().expand(100, 1, 64, 64)
>>>
>>> data = MyData(images, mask, label=label, batch_size=[100])
>>>
>>> print(data.select_label(1))
MyData(
    image=Tensor(torch.Size([11, 3, 64, 64]), dtype=torch.float32),
    label=Tensor(torch.Size([11]), dtype=torch.int64),
    mask=Tensor(torch.Size([11, 1, 64, 64]), dtype=torch.bool),
    batch_size=torch.Size([11]),
    device=None,
    is_shared=False)
>>> print(data.mask_image().shape)
torch.Size([100, 6117])
>>> print(data.reshape(10, 10))
MyData(
    image=Tensor(torch.Size([10, 10, 3, 64, 64]), dtype=torch.float32),
    label=Tensor(torch.Size([10, 10]), dtype=torch.int64),
    mask=Tensor(torch.Size([10, 10, 1, 64, 64]), dtype=torch.bool),
    batch_size=torch.Size([10, 10]),
    device=None,
    is_shared=False)

As this example shows, one can write a specific data structures with dedicated methods while still enjoying the TensorDict artifacts such as shape operations (e.g. reshape or permutations), data manipulation (indexing, cat and stack) or calling arbitrary functions through the apply method (and many more).

Tensorclasses support nesting and, in fact, all the TensorDict features.

Installation

With Pip:

To install the latest stable version of tensordict, simply run

pip install tensordict

This will work with Python 3.7 and upward as well as PyTorch 1.12 and upward.

To enjoy the latest features, one can use

pip install tensordict-nightly

With Conda:

Install tensordict from conda-forge channel.

conda install -c conda-forge tensordict

Citation

If you're using TensorDict, please refer to this BibTeX entry to cite this work:

@misc{bou2023torchrl,
      title={TorchRL: A data-driven decision-making library for PyTorch}, 
      author={Albert Bou and Matteo Bettini and Sebastian Dittert and Vikash Kumar and Shagun Sodhani and Xiaomeng Yang and Gianni De Fabritiis and Vincent Moens},
      year={2023},
      eprint={2306.00577},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Disclaimer

TensorDict is at the beta-stage, meaning that there may be bc-breaking changes introduced, but they should come with a warranty. Hopefully these should not happen too often, as the current roadmap mostly involves adding new features and building compatibility with the broader PyTorch ecosystem.

License

TensorDict is licensed under the MIT License. See LICENSE for details.

Project details


Release history Release notifications | RSS feed

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

tensordict_nightly-2023.8.31-cp311-cp311-win_amd64.whl (217.9 kB view details)

Uploaded CPython 3.11 Windows x86-64

tensordict_nightly-2023.8.31-cp311-cp311-macosx_10_9_universal2.whl (277.6 kB view details)

Uploaded CPython 3.11 macOS 10.9+ universal2 (ARM64, x86-64)

tensordict_nightly-2023.8.31-cp310-cp310-win_amd64.whl (217.4 kB view details)

Uploaded CPython 3.10 Windows x86-64

tensordict_nightly-2023.8.31-cp310-cp310-macosx_10_15_x86_64.whl (219.5 kB view details)

Uploaded CPython 3.10 macOS 10.15+ x86-64

tensordict_nightly-2023.8.31-cp39-cp39-win_amd64.whl (217.4 kB view details)

Uploaded CPython 3.9 Windows x86-64

tensordict_nightly-2023.8.31-cp39-cp39-macosx_11_0_x86_64.whl (219.6 kB view details)

Uploaded CPython 3.9 macOS 11.0+ x86-64

tensordict_nightly-2023.8.31-cp38-cp38-win_amd64.whl (217.4 kB view details)

Uploaded CPython 3.8 Windows x86-64

tensordict_nightly-2023.8.31-cp38-cp38-macosx_11_0_x86_64.whl (219.4 kB view details)

Uploaded CPython 3.8 macOS 11.0+ x86-64

File details

Details for the file tensordict_nightly-2023.8.31-cp311-cp311-win_amd64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 841a19e0422638555f115fc3abb69cd048cf06ff2a49c6d541884fc3c3516b20
MD5 8ba3280329d0fb1e3c2b86404154e155
BLAKE2b-256 83d76a1262d5824c7f6fd7de38625626a1467a92068e235d0ef56a5e06d4acf8

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp311-cp311-manylinux1_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 7b3d3b0ace9bdf89a88946acd3fc53621a7ca681b4994641dfe85c1d6af7faa7
MD5 4b516615e5ee85e2bb8423f286405cbd
BLAKE2b-256 8c406383aa4e0b82602d998eed20004898fde5e2ab30d7936b2aafc36c8f66f1

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp311-cp311-macosx_10_9_universal2.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp311-cp311-macosx_10_9_universal2.whl
Algorithm Hash digest
SHA256 f647cb33dba0c76f59e34f93fb49bc1bf5ce2c6a9dc37a01d84080979c85d15c
MD5 1e29167890874e83c4fc84b5f098a77a
BLAKE2b-256 4f59b486d2f9e8e2caf2395a4e1029e1adeafb899c64fd7cfe51cfdabd07baa7

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp310-cp310-win_amd64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 04dfbc45081d64774da86e28e654df5f90d1d1cb633b32daf067e039a82cdbb7
MD5 babb602dd2e1c40f218dd09d5b7e1d36
BLAKE2b-256 095ab7d616a530d4a534b45ff833f01870ff0faef4f062264c3cc821821ed66d

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp310-cp310-manylinux1_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 603ca874e34c0a25e99eba91c6e4f00dc32ca7c429ae13d703f85867a9090c22
MD5 4f84b5d11ba9353eb02c930689523ae9
BLAKE2b-256 3889a0fc4a7f418385ada12eec1fe3730f9c9a55d54e67500c28fb600157c1d3

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp310-cp310-macosx_10_15_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp310-cp310-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 1d82849cd7f41338790245d9fbbe440664b3dfe6e043ff69dc135d1f46f559d0
MD5 d53e8e993ce4d5692404b58ebbbfc0fd
BLAKE2b-256 9ad609c8ff94acdbe1897ab83232cb3afb49f498403e4ae75a0421dfd9e6cc21

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp39-cp39-win_amd64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp39-cp39-win_amd64.whl
Algorithm Hash digest
SHA256 28abeaa59c97a6e38e60495708571806a2837a13a3e5793c93452aceb0b2e570
MD5 ba63556367d9fa0a867229e3553467b0
BLAKE2b-256 ed114daac13dcd4c1d64e81c3029809bb2635da7369c6ee66209e36e438362b2

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp39-cp39-manylinux1_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp39-cp39-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 60acb311a49457f09867d2fdc193b5655a3d517cbdb1755e1f3d247fa8a30d1f
MD5 d17882fa46612125fd584007f899d5d0
BLAKE2b-256 ee936f95246fc975f62c1ce03ada69668103d4c92c0062db5481a21941252842

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp39-cp39-macosx_11_0_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp39-cp39-macosx_11_0_x86_64.whl
Algorithm Hash digest
SHA256 96eea0f4bdea58a04754368107b267e0c26fe0991d4f1f4818c2200ca53ea6a1
MD5 c12e7b46b68380c70ebfdb649c766853
BLAKE2b-256 7ae3c9dab514f63be87e7802ee449cb62a02f94feff60535c22d384f916ab5e0

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp38-cp38-win_amd64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp38-cp38-win_amd64.whl
Algorithm Hash digest
SHA256 054b615eda3c4b11c29b793d97aab315f0adacfeb17735eff45e35b1e9d982b4
MD5 f6701ac2608b9c4a3118e66b633d2ebc
BLAKE2b-256 9c9aa4d81785e2a46c396a3b35d2ce6bf650ba39a2ad9d798e066b55e04e03ad

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp38-cp38-manylinux1_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp38-cp38-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 a5e7781c2fab7229c555d73e6924aca693969a2d7404b93ba5cf2a44cb6ae1f9
MD5 ef538ca9b3da233ef6bc5bc5737b9e84
BLAKE2b-256 59e5cffa087a5e20bd8be7e32a3a56a307ad8b42fbe9235df262034505cf4e36

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2023.8.31-cp38-cp38-macosx_11_0_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2023.8.31-cp38-cp38-macosx_11_0_x86_64.whl
Algorithm Hash digest
SHA256 b36189cef792be68d46349e00d2098a26c675c544a18757118197987d92eed06
MD5 8dce92ed2d00a602474881ada4d60e80
BLAKE2b-256 7de101c6fcc7a88faab7eb0332bd90103f9ccd593b8620eca66443f11f50e609

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page