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.11.23-cp311-cp311-win_amd64.whl (241.8 kB view details)

Uploaded CPython 3.11 Windows x86-64

tensordict_nightly-2023.11.23-cp311-cp311-macosx_10_9_universal2.whl (301.1 kB view details)

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

tensordict_nightly-2023.11.23-cp310-cp310-win_amd64.whl (241.3 kB view details)

Uploaded CPython 3.10 Windows x86-64

tensordict_nightly-2023.11.23-cp310-cp310-macosx_10_15_x86_64.whl (242.9 kB view details)

Uploaded CPython 3.10 macOS 10.15+ x86-64

tensordict_nightly-2023.11.23-cp39-cp39-win_amd64.whl (240.9 kB view details)

Uploaded CPython 3.9 Windows x86-64

tensordict_nightly-2023.11.23-cp39-cp39-macosx_11_0_x86_64.whl (243.1 kB view details)

Uploaded CPython 3.9 macOS 11.0+ x86-64

tensordict_nightly-2023.11.23-cp38-cp38-win_amd64.whl (241.2 kB view details)

Uploaded CPython 3.8 Windows x86-64

tensordict_nightly-2023.11.23-cp38-cp38-macosx_11_0_x86_64.whl (242.9 kB view details)

Uploaded CPython 3.8 macOS 11.0+ x86-64

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 4c62a73450d0642d5b8b41a65af4c7d3aedd3dcb0e418340c5911796d8af7003
MD5 3b11e19ce51fd25cc052011a25360d11
BLAKE2b-256 1161bfcbc883cb5d681f978272f571340f2c004ef2b919d8bf7244754cbecfc8

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 7bbf313ad04a32940313d798f631cb0cdebbdd9811b25bd86237466c35d7834c
MD5 5ac1be2bfb1d8e77bc13cce1d438bcc9
BLAKE2b-256 78a21cee331db9e4f5890953e59130279a94aabd658d2fb96f58c1d3dc05d095

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp311-cp311-macosx_10_9_universal2.whl
Algorithm Hash digest
SHA256 cdb95705e42d00f0c642b57a6e5a3d73e20dba8c7ff4395ce035accd56fb90e2
MD5 167101aa8ef0ed179c2c2ce847b3cb14
BLAKE2b-256 2ca4c82d2467f020ac7e37ae67ebb0d5c5c5edba8ee931e97567fd017a2cf404

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 7a61c081a2af4dc549beb07d5e99b88be9f138497f010d59353b68422c555e9c
MD5 b10515f0a4f4ac141f26b0f02c5a531d
BLAKE2b-256 189cbba2530c7824c5df319974016daa10b5160f4cf12fab2e658583467b3611

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 2b96746f88b433292174d6f2711b1a1bdc3934d78c525048a61cd9bd691122f1
MD5 9dfdeb6ca2c5239863ce39884bc61ab8
BLAKE2b-256 521b4e4d73ea3c7a7d34b5500f021c76acb7585c7f36e7486af9065ce08505dc

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp310-cp310-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 ea7d3728464b543a9abd5bd9ac163d4380b23078730b8c04633e76925d513e8c
MD5 2e9714e108434e00febec9c9413eef49
BLAKE2b-256 4836fbcc676b7b4aabf8a0dea2b51bd00f8a37c485e8a8e160a224b9d2320733

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp39-cp39-win_amd64.whl
Algorithm Hash digest
SHA256 dfaf981bd82612d0ce3f9f416ebbd67207530a16ef3212fefc92e926d4af4ac9
MD5 6354dc03e10dd459b57511f31977273c
BLAKE2b-256 1b4f1cdf6c89446ac6bec1a1d6beb2c0d65f2dec003727ec8f2ce3b457934b48

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp39-cp39-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 3854f663e9f76210eb8533f91b2e991ef0a3bfdf0e7d0a9aad7f0deca16a313a
MD5 a7ae118f6c2a76717841f9380f05915f
BLAKE2b-256 ffe3d8e8544de4e02a4123fcb23f2884e8f77e7301ff654956f4bb038bd717ce

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp39-cp39-macosx_11_0_x86_64.whl
Algorithm Hash digest
SHA256 a976491e5af1bc2335f0d9517420dac6381f38d37231f659baeaeea52852225b
MD5 23a077fb95bb97dccfe15646ad8a4726
BLAKE2b-256 fa6aeaed96a72ab57e770967424d4b8ada4f265c330271d6785c303d0c780217

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp38-cp38-win_amd64.whl
Algorithm Hash digest
SHA256 3764ac060359696b70957686afa2be6220de030386535c3a88de440584608d67
MD5 35b1d5b11e2aff55620e14dfc90aee11
BLAKE2b-256 c37a1add9fdef233d10ab1c050aecbf0d5bfa0de52bd616bc3271fe2e0bfeb5a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp38-cp38-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 6e14ae6e532bc2ecce410aa9ed5409fca35dd2dd7ade703a305f8238e9ac2c3c
MD5 cf4974778d41be23da16a884ae8ac179
BLAKE2b-256 24327f8e4bc6c3befe68a0297c045464a4c6a8314ddcc2a8ad105597de107769

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2023.11.23-cp38-cp38-macosx_11_0_x86_64.whl
Algorithm Hash digest
SHA256 a3d2cec8b8a140a22e9afe798ca9006d77d8388b1762c7209bffbe4f6de3a172
MD5 fc481e63b336ead0b18b2a72aaadd4e3
BLAKE2b-256 e9f6c56657cb403d9bd1ef4bc3df11bcc2cfa0eda4fdc2f7c61196c97b023065

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