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

Uploaded CPython 3.11 Windows x86-64

tensordict_nightly-2024.1.13-cp311-cp311-macosx_10_9_universal2.whl (315.0 kB view details)

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

tensordict_nightly-2024.1.13-cp310-cp310-win_amd64.whl (255.3 kB view details)

Uploaded CPython 3.10 Windows x86-64

tensordict_nightly-2024.1.13-cp310-cp310-macosx_10_15_x86_64.whl (256.8 kB view details)

Uploaded CPython 3.10 macOS 10.15+ x86-64

tensordict_nightly-2024.1.13-cp39-cp39-win_amd64.whl (254.9 kB view details)

Uploaded CPython 3.9 Windows x86-64

tensordict_nightly-2024.1.13-cp39-cp39-macosx_11_0_x86_64.whl (257.0 kB view details)

Uploaded CPython 3.9 macOS 11.0+ x86-64

tensordict_nightly-2024.1.13-cp38-cp38-win_amd64.whl (255.2 kB view details)

Uploaded CPython 3.8 Windows x86-64

tensordict_nightly-2024.1.13-cp38-cp38-macosx_11_0_x86_64.whl (256.7 kB view details)

Uploaded CPython 3.8 macOS 11.0+ x86-64

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 4ab58541a367d31c15103d1f04f62c94c3ad1a8308dc710d4cd9b4d843c52aac
MD5 f9f0b238bc954b9420b2a16546b0ae78
BLAKE2b-256 4fab0fa1d9671e4314bf73164a1ebd5f29edc3a19aeb41019dc614cc96401816

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 4cc4e514d818c30629c3247a266c407b813e1ed830f36d8d145f237e8eb9f86c
MD5 f0948965ed81c8ee44871bc3101f43fd
BLAKE2b-256 d95ae7a55492ecafdb43b407a9f89771959ba39b30942a5dce68f4768cd5df00

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp311-cp311-macosx_10_9_universal2.whl
Algorithm Hash digest
SHA256 e1675f44c519044dd3466cf590a542dac12fe2501b59c3a6026d631101ba0869
MD5 7cdb8e757fcca2d90353ba2f1cff23cb
BLAKE2b-256 14ab93a44f176fabb731f8176f667666c1da588c77bce02b4a91c2442da6d754

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 3a1d4362cc2da31c167016b5cfaa144b73d5861ed900150863611d4d2b832520
MD5 93af6dbba94aca45c602c93651e85d66
BLAKE2b-256 0005c61f336b921912dae058bffca276caa5f91367e59d32ff95ddf4682a8038

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 ccd98c9856ff67c0807b75dca27058dac0e39db4a3ead144db111258d442bb6c
MD5 4771e35d65bc612b6c4c972947590a8b
BLAKE2b-256 c7a6361b2f87e95095fe185b9c4cebbf6adf61b89513f1ed54303d8d5e131c86

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp310-cp310-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 c298c2af36b5d6965c645cdf403f942d40c669e8268e7979f8f8884d82dbf827
MD5 667e5ad09d450ed6217a4d148ca50d78
BLAKE2b-256 4d1295e3d7490ca216f2b81d7ceb03ed445bad1325c4c2e0d205cdfac430b451

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp39-cp39-win_amd64.whl
Algorithm Hash digest
SHA256 27dc7362721759fb1cd5a8143a51bac852cc551bb6da01fc2a269279f3d7f5cb
MD5 7199807831485d3f8e554724a7f7cc2a
BLAKE2b-256 5ce6439ade55a38fa21b2bdbe30f11e258e2cef880e8d7e026ea323f5b985944

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp39-cp39-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 fe953da25b3545cdb185fdcc4969add26910ea41227bd130bbc92413de1b9911
MD5 ffcc3b911d4e206230d20a7c525ef591
BLAKE2b-256 112ad8a4533fadfd44f66d3fe17e198fb3886e1eb6018a92bf8bc509bf791900

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp39-cp39-macosx_11_0_x86_64.whl
Algorithm Hash digest
SHA256 cae6df874c266195ead57c5ed594340fa5728ca3cb735602f63bbb800ce153ac
MD5 9401d242c9b58b619b8fc4149b80bad3
BLAKE2b-256 36de0abf711a1396d9f76f64adbcb9b5fd089fafa2b4b5c6dbbf38373cbb75dd

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp38-cp38-win_amd64.whl
Algorithm Hash digest
SHA256 f0ef62f972bff450712daf67b83826756c659b55a5f780aa8ad12e9d2aad8a31
MD5 55a1b1dc9aba6ce3e977950e3d311441
BLAKE2b-256 ac244811a4ebf077f13c50d15affa15104604973835c9b9654583158ecffaec7

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp38-cp38-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 3ead85ce11cd3d5de120a2b7b3cd303f551195b5261b5302d94ffc4297080a1c
MD5 3c9c36ebf2a87f9762bdf48dad166ddb
BLAKE2b-256 f0a1a4dd020e9a5d9e714fcf51d1041a966fa2f2087806b4e850f2b8fbea8613

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2024.1.13-cp38-cp38-macosx_11_0_x86_64.whl
Algorithm Hash digest
SHA256 d95d037732d6df3b4402fe258ea93d45057194e7068c9dd6d7b115bd40011f65
MD5 b96841442a458306ece31945758aeed3
BLAKE2b-256 3919a509f830be2549ae1023359411efcda79487f116074f5b446e85ae416922

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