<!--- BADGES: START --->
<!---
[](https://pytorch.github.io/tensordict/)
--->
[][#docs-package]
[](https://discord.gg/tz3TgTAe3D)
[][#docs-package-benchmark]
[](https://www.python.org/downloads/)
[][#github-license]
<a href="https://pypi.org/project/tensordict"><img src="https://img.shields.io/pypi/v/tensordict" alt="pypi version"></a>
<a href="https://pypi.org/project/tensordict-nightly"><img src="https://img.shields.io/pypi/v/tensordict-nightly?label=nightly" alt="pypi nightly version"></a>
[][#pepy-package]
[)][#pepy-package-nightly]
[][#codecov-package]
[][#circleci-package]
[][#conda-forge-package]
[][#conda-forge-package]
[#docs-package]: https://pytorch.github.io/tensordict/
[#docs-package-benchmark]: https://pytorch.github.io/tensordict/dev/bench/
[#github-license]: https://github.com/pytorch/tensordict/blob/main/LICENSE
[#pepy-package]: https://pepy.tech/project/tensordict
[#pepy-package-nightly]: https://pepy.tech/project/tensordict-nightly
[#codecov-package]: https://codecov.io/gh/pytorch/tensordict
[#circleci-package]: https://circleci.com/gh/pytorch/tensordict
[#conda-forge-package]: https://anaconda.org/conda-forge/tensordict
<!--- BADGES: END --->
# ���� TensorDict
TensorDict is a dictionary-like class that inherits properties from tensors, making it easy to work with collections of
tensors in PyTorch. It provides a simple and intuitive way to manipulate and process tensors, allowing you to focus on
building and training your models.
[**Key Features**](#key-features) |
[**Examples**](#examples) |
[**Installation**](#installation) |
[**Citation**](#citation) |
[**License**](#license)
## Key Features
TensorDict makes your code-bases more _readable_, _compact_, _modular_ and _fast_.
It abstracts away tailored operations, making your code less error-prone as it takes care of
dispatching the operation on the leaves for you.
The key features are:
- ���� **Composability**: `TensorDict` generalizes `torch.Tensor` operations to collection of tensors.
- ������ **Speed**: asynchronous transfer to device, fast node-to-node communication through `consolidate`, compatible with `torch.compile`.
- ������ **Shape operations**: Perform tensor-like operations on TensorDict instances, such as indexing, slicing or
concatenation.
- ���� **Distributed / multiprocessed capabilities**: Easily distribute TensorDict instances across multiple workers,
devices and machines.
- ���� **Serialization** and memory-mapping
- �� **Functional programming** and compatibility with `torch.vmap`
- ���� **Nesting**: Nest TensorDict instances to create hierarchical structures.
- ��� **Lazy preallocation**: Preallocate memory for TensorDict instances without initializing the tensors.
- ���� **Specialized dataclass** for torch.Tensor ([`@tensorclass`](#tensorclass))
## Examples
This section presents a couple of stand-out applications of the library.
Check our [**Getting Started**](GETTING_STARTED.md) guide for an overview of TensorDict's features!
### Fast copy on device
`TensorDict` optimizes transfers from/to device to make them safe and fast.
By default, data transfers will be made asynchronously and synchronizations will be called whenever needed.
```python
# Fast and safe asynchronous copy to 'cuda'
td_cuda = TensorDict(**dict_of_tensor, device="cuda")
# Fast and safe asynchronous copy to 'cpu'
td_cpu = td_cuda.to("cpu")
# Force synchronous copy
td_cpu = td_cuda.to("cpu", non_blocking=False)
```
### Coding an optimizer
For instance, using `TensorDict` you can code the Adam optimizer as you would for a single `torch.Tensor` and apply
that to a `TensorDict` input as well. On `cuda`, these operations will rely on fused kernels, making it very fast to
execute:
```python
class Adam:
def __init__(self, weights: TensorDict, alpha: float=1e-3,
beta1: float=0.9, beta2: float=0.999,
eps: float = 1e-6):
# Lock for efficiency
weights = weights.lock_()
self.weights = weights
self.t = 0
self._mu = weights.data.clone()
self._sigma = weights.data.mul(0.0)
self.beta1 = beta1
self.beta2 = beta2
self.alpha = alpha
self.eps = eps
def step(self):
self._mu.mul_(self.beta1).add_(self.weights.grad, 1 - self.beta1)
self._sigma.mul_(self.beta2).add_(self.weights.grad.pow(2), 1 - self.beta2)
self.t += 1
mu = self._mu.div_(1-self.beta1**self.t)
sigma = self._sigma.div_(1 - self.beta2 ** self.t)
self.weights.data.add_(mu.div_(sigma.sqrt_().add_(self.eps)).mul_(-self.alpha))
```
### Training a model
Using tensordict primitives, most supervised training loops can be rewritten in a generic way:
```python
for i, data in enumerate(dataset):
# the model reads and writes tensordicts
data = model(data)
loss = loss_module(data)
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
**With Pip**:
To install the latest stable version of tensordict, simply run
```bash
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
```bash
pip install tensordict-nightly
```
**With Conda**:
Install `tensordict` from `conda-forge` channel.
```sh
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](LICENSE) for details.
Raw data
{
"_id": null,
"home_page": "https://github.com/pytorch/tensordict",
"name": "tensordict",
"maintainer": null,
"docs_url": null,
"requires_python": null,
"maintainer_email": null,
"keywords": null,
"author": "tensordict contributors",
"author_email": "vmoens@fb.com",
"download_url": null,
"platform": null,
"description": "<!--- BADGES: START --->\n<!---\n[](https://pytorch.github.io/tensordict/)\n--->\n[][#docs-package]\n[](https://discord.gg/tz3TgTAe3D)\n[][#docs-package-benchmark]\n[](https://www.python.org/downloads/)\n[][#github-license]\n<a href=\"https://pypi.org/project/tensordict\"><img src=\"https://img.shields.io/pypi/v/tensordict\" alt=\"pypi version\"></a>\n<a href=\"https://pypi.org/project/tensordict-nightly\"><img src=\"https://img.shields.io/pypi/v/tensordict-nightly?label=nightly\" alt=\"pypi nightly version\"></a>\n[][#pepy-package]\n[)][#pepy-package-nightly]\n[][#codecov-package]\n[][#circleci-package]\n[][#conda-forge-package]\n[][#conda-forge-package]\n\n[#docs-package]: https://pytorch.github.io/tensordict/\n[#docs-package-benchmark]: https://pytorch.github.io/tensordict/dev/bench/\n[#github-license]: https://github.com/pytorch/tensordict/blob/main/LICENSE\n[#pepy-package]: https://pepy.tech/project/tensordict\n[#pepy-package-nightly]: https://pepy.tech/project/tensordict-nightly\n[#codecov-package]: https://codecov.io/gh/pytorch/tensordict\n[#circleci-package]: https://circleci.com/gh/pytorch/tensordict\n[#conda-forge-package]: https://anaconda.org/conda-forge/tensordict\n\n<!--- BADGES: END --->\n\n# \ufffd\ufffd\ufffd\ufffd TensorDict\n\nTensorDict is a dictionary-like class that inherits properties from tensors, making it easy to work with collections of\ntensors in PyTorch. It provides a simple and intuitive way to manipulate and process tensors, allowing you to focus on\nbuilding and training your models.\n\n[**Key Features**](#key-features) |\n[**Examples**](#examples) |\n[**Installation**](#installation) |\n[**Citation**](#citation) |\n[**License**](#license)\n\n## Key Features\nTensorDict makes your code-bases more _readable_, _compact_, _modular_ and _fast_.\nIt abstracts away tailored operations, making your code less error-prone as it takes care of\ndispatching the operation on the leaves for you.\n\nThe key features are:\n\n- \ufffd\ufffd\ufffd\ufffd **Composability**: `TensorDict` generalizes `torch.Tensor` operations to collection of tensors.\n- \ufffd\ufffd\ufffd\ufffd\ufffd\ufffd **Speed**: asynchronous transfer to device, fast node-to-node communication through `consolidate`, compatible with `torch.compile`.\n- \ufffd\ufffd\ufffd\ufffd\ufffd\ufffd **Shape operations**: Perform tensor-like operations on TensorDict instances, such as indexing, slicing or\n concatenation.\n- \ufffd\ufffd\ufffd\ufffd **Distributed / multiprocessed capabilities**: Easily distribute TensorDict instances across multiple workers,\n devices and machines.\n- \ufffd\ufffd\ufffd\ufffd **Serialization** and memory-mapping\n- \ufffd\ufffd **Functional programming** and compatibility with `torch.vmap`\n- \ufffd\ufffd\ufffd\ufffd **Nesting**: Nest TensorDict instances to create hierarchical structures.\n- \ufffd\ufffd\ufffd **Lazy preallocation**: Preallocate memory for TensorDict instances without initializing the tensors.\n- \ufffd\ufffd\ufffd\ufffd **Specialized dataclass** for torch.Tensor ([`@tensorclass`](#tensorclass))\n\n\n\n## Examples\n\nThis section presents a couple of stand-out applications of the library.\nCheck our [**Getting Started**](GETTING_STARTED.md) guide for an overview of TensorDict's features!\n\n### Fast copy on device\n`TensorDict` optimizes transfers from/to device to make them safe and fast.\nBy default, data transfers will be made asynchronously and synchronizations will be called whenever needed.\n```python\n# Fast and safe asynchronous copy to 'cuda'\ntd_cuda = TensorDict(**dict_of_tensor, device=\"cuda\")\n# Fast and safe asynchronous copy to 'cpu'\ntd_cpu = td_cuda.to(\"cpu\")\n# Force synchronous copy\ntd_cpu = td_cuda.to(\"cpu\", non_blocking=False)\n```\n\n### Coding an optimizer\nFor instance, using `TensorDict` you can code the Adam optimizer as you would for a single `torch.Tensor` and apply\nthat to a `TensorDict` input as well. On `cuda`, these operations will rely on fused kernels, making it very fast to\nexecute:\n```python\nclass Adam:\n def __init__(self, weights: TensorDict, alpha: float=1e-3,\n beta1: float=0.9, beta2: float=0.999,\n eps: float = 1e-6):\n # Lock for efficiency\n weights = weights.lock_()\n self.weights = weights\n self.t = 0\n\n self._mu = weights.data.clone()\n self._sigma = weights.data.mul(0.0)\n self.beta1 = beta1\n self.beta2 = beta2\n self.alpha = alpha\n self.eps = eps\n\n def step(self):\n self._mu.mul_(self.beta1).add_(self.weights.grad, 1 - self.beta1)\n self._sigma.mul_(self.beta2).add_(self.weights.grad.pow(2), 1 - self.beta2)\n self.t += 1\n mu = self._mu.div_(1-self.beta1**self.t)\n sigma = self._sigma.div_(1 - self.beta2 ** self.t)\n self.weights.data.add_(mu.div_(sigma.sqrt_().add_(self.eps)).mul_(-self.alpha))\n```\n\n### Training a model\n\nUsing tensordict primitives, most supervised training loops can be rewritten in a generic way:\n```python\nfor i, data in enumerate(dataset):\n # the model reads and writes tensordicts\n data = model(data)\n loss = loss_module(data)\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n```\n\nWith this level of abstraction, one can recycle a training loop for highly heterogeneous task.\nEach individual step of the training loop (data collection and transform, model prediction, loss computation etc.)\ncan be tailored to the use case at hand without impacting the others.\nFor instance, the above example can be easily used across classification and segmentation tasks, among many others.\n\n## Installation\n\n**With Pip**:\n\nTo install the latest stable version of tensordict, simply run\n\n```bash\npip install tensordict\n```\n\nThis will work with Python 3.7 and upward as well as PyTorch 1.12 and upward.\n\nTo enjoy the latest features, one can use\n\n```bash\npip install tensordict-nightly\n```\n\n**With Conda**:\n\nInstall `tensordict` from `conda-forge` channel.\n\n```sh\nconda install -c conda-forge tensordict\n```\n\n## Citation\n\nIf you're using TensorDict, please refer to this BibTeX entry to cite this work:\n```\n@misc{bou2023torchrl,\n title={TorchRL: A data-driven decision-making library for PyTorch},\n 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},\n year={2023},\n eprint={2306.00577},\n archivePrefix={arXiv},\n primaryClass={cs.LG}\n}\n```\n\n## Disclaimer\n\nTensorDict is at the *beta*-stage, meaning that there may be bc-breaking changes introduced, but\nthey should come with a warranty.\nHopefully these should not happen too often, as the current roadmap mostly\ninvolves adding new features and building compatibility with the broader\nPyTorch ecosystem.\n\n## License\n\nTensorDict is licensed under the MIT License. See [LICENSE](LICENSE) for details.\n",
"bugtrack_url": null,
"license": "BSD",
"summary": null,
"version": "0.6.2",
"project_urls": {
"Homepage": "https://github.com/pytorch/tensordict"
},
"split_keywords": [],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "a9b140afa9e2e11414b8def6fedb044eb77ebdfb17ab02b1d925339e2dce3ada",
"md5": "54bc95c6d7cebcc76425c71f9f491bfa",
"sha256": "a00aea0dd7d38026b7e4721d3e4c9759fed99829d8f3d337f4f5174c42138c87"
},
"downloads": -1,
"filename": "tensordict-0.6.2-cp310-cp310-macosx_11_0_arm64.whl",
"has_sig": false,
"md5_digest": "54bc95c6d7cebcc76425c71f9f491bfa",
"packagetype": "bdist_wheel",
"python_version": "cp310",
"requires_python": null,
"size": 668393,
"upload_time": "2024-11-14T15:01:26",
"upload_time_iso_8601": "2024-11-14T15:01:26.052980Z",
"url": "https://files.pythonhosted.org/packages/a9/b1/40afa9e2e11414b8def6fedb044eb77ebdfb17ab02b1d925339e2dce3ada/tensordict-0.6.2-cp310-cp310-macosx_11_0_arm64.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "96d0ecdb847d7190cca1e6dcdfbfd374fad82937b890bd4609385cdf3b4d5592",
"md5": "b7285d44ec759fe97abeb98a458ef1be",
"sha256": "24a505926cca17401411b9ed7377013c57dd49d579f350ac2e002d701a7f6cfd"
},
"downloads": -1,
"filename": "tensordict-0.6.2-cp310-cp310-manylinux1_x86_64.whl",
"has_sig": false,
"md5_digest": "b7285d44ec759fe97abeb98a458ef1be",
"packagetype": "bdist_wheel",
"python_version": "cp310",
"requires_python": null,
"size": 359852,
"upload_time": "2024-11-14T15:02:31",
"upload_time_iso_8601": "2024-11-14T15:02:31.321606Z",
"url": "https://files.pythonhosted.org/packages/96/d0/ecdb847d7190cca1e6dcdfbfd374fad82937b890bd4609385cdf3b4d5592/tensordict-0.6.2-cp310-cp310-manylinux1_x86_64.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "fdc4ddbd7109fef983a726cc135bdd748a7b3747cd3e77dc300ab4cc577a8d64",
"md5": "108987cae370b49b804c53c6c8d784f7",
"sha256": "e5132b369ac9e84eb0e9e28d1954f6d1370e06927ab3c502e02a6d35a1ed629a"
},
"downloads": -1,
"filename": "tensordict-0.6.2-cp311-cp311-macosx_11_0_arm64.whl",
"has_sig": false,
"md5_digest": "108987cae370b49b804c53c6c8d784f7",
"packagetype": "bdist_wheel",
"python_version": "cp311",
"requires_python": null,
"size": 669700,
"upload_time": "2024-11-14T15:02:33",
"upload_time_iso_8601": "2024-11-14T15:02:33.772682Z",
"url": "https://files.pythonhosted.org/packages/fd/c4/ddbd7109fef983a726cc135bdd748a7b3747cd3e77dc300ab4cc577a8d64/tensordict-0.6.2-cp311-cp311-macosx_11_0_arm64.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "2904a58ed454ad7f8f8a65f6ebe2df279c4d6921234cad3603318cf8a150b0f3",
"md5": "ddf7c2f6477656c0efac22bec0380e7e",
"sha256": "3e036ae56532b6308a9e95de29db97f1ec9b83f569c716a602d7d48dc0b47ef5"
},
"downloads": -1,
"filename": "tensordict-0.6.2-cp311-cp311-manylinux1_x86_64.whl",
"has_sig": false,
"md5_digest": "ddf7c2f6477656c0efac22bec0380e7e",
"packagetype": "bdist_wheel",
"python_version": "cp311",
"requires_python": null,
"size": 360364,
"upload_time": "2024-11-14T15:02:35",
"upload_time_iso_8601": "2024-11-14T15:02:35.838993Z",
"url": "https://files.pythonhosted.org/packages/29/04/a58ed454ad7f8f8a65f6ebe2df279c4d6921234cad3603318cf8a150b0f3/tensordict-0.6.2-cp311-cp311-manylinux1_x86_64.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "d217a2bdb51d5d1b2f06debdf8f60c159cf328c778b6cc8103940f2079e2e767",
"md5": "9eb7cba4934ad21bae1ec7d7a133f239",
"sha256": "529515ffcb0c6ab5d1c0c85798defe8972e5390412c356c79b87f4073e48b537"
},
"downloads": -1,
"filename": "tensordict-0.6.2-cp312-cp312-macosx_11_0_arm64.whl",
"has_sig": false,
"md5_digest": "9eb7cba4934ad21bae1ec7d7a133f239",
"packagetype": "bdist_wheel",
"python_version": "cp312",
"requires_python": null,
"size": 668892,
"upload_time": "2024-11-14T15:02:37",
"upload_time_iso_8601": "2024-11-14T15:02:37.036194Z",
"url": "https://files.pythonhosted.org/packages/d2/17/a2bdb51d5d1b2f06debdf8f60c159cf328c778b6cc8103940f2079e2e767/tensordict-0.6.2-cp312-cp312-macosx_11_0_arm64.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "64aa2b562740c8be079aba4bb9eb37cb94ed042d65e48f739577eb70eee617c2",
"md5": "c3e4a9d7cca76a5bc4cfd7619f0638b3",
"sha256": "6ce86d8f9ed53d1c893e0064729a799eb06371902e2d5a2f3c83c6af010cf3c2"
},
"downloads": -1,
"filename": "tensordict-0.6.2-cp312-cp312-manylinux1_x86_64.whl",
"has_sig": false,
"md5_digest": "c3e4a9d7cca76a5bc4cfd7619f0638b3",
"packagetype": "bdist_wheel",
"python_version": "cp312",
"requires_python": null,
"size": 360647,
"upload_time": "2024-11-14T15:02:38",
"upload_time_iso_8601": "2024-11-14T15:02:38.475998Z",
"url": "https://files.pythonhosted.org/packages/64/aa/2b562740c8be079aba4bb9eb37cb94ed042d65e48f739577eb70eee617c2/tensordict-0.6.2-cp312-cp312-manylinux1_x86_64.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "3bca58c7010737dc0dfea44824e1691ac679c69e93d824c907ab1145c2c0f5e9",
"md5": "5933e8950c9626b09c9c91388473e5ba",
"sha256": "1bfb1d41877fc19b709a7ad20c3b36f19224e5013263989de8e994bd98559fda"
},
"downloads": -1,
"filename": "tensordict-0.6.2-cp39-cp39-macosx_11_0_arm64.whl",
"has_sig": false,
"md5_digest": "5933e8950c9626b09c9c91388473e5ba",
"packagetype": "bdist_wheel",
"python_version": "cp39",
"requires_python": null,
"size": 668536,
"upload_time": "2024-11-14T15:02:40",
"upload_time_iso_8601": "2024-11-14T15:02:40.196294Z",
"url": "https://files.pythonhosted.org/packages/3b/ca/58c7010737dc0dfea44824e1691ac679c69e93d824c907ab1145c2c0f5e9/tensordict-0.6.2-cp39-cp39-macosx_11_0_arm64.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "2ca8031063d078e09fff8f3bed5d07707e50912fb9c4cff4f65c032eb165b432",
"md5": "403b25bdcf6de9d7ab1d4a45a2ffd9f6",
"sha256": "490371099494cdb3250b77d64e5ea63da8c3715a612ffab03e5f9a791161bffd"
},
"downloads": -1,
"filename": "tensordict-0.6.2-cp39-cp39-manylinux1_x86_64.whl",
"has_sig": false,
"md5_digest": "403b25bdcf6de9d7ab1d4a45a2ffd9f6",
"packagetype": "bdist_wheel",
"python_version": "cp39",
"requires_python": null,
"size": 359877,
"upload_time": "2024-11-14T15:02:41",
"upload_time_iso_8601": "2024-11-14T15:02:41.985814Z",
"url": "https://files.pythonhosted.org/packages/2c/a8/031063d078e09fff8f3bed5d07707e50912fb9c4cff4f65c032eb165b432/tensordict-0.6.2-cp39-cp39-manylinux1_x86_64.whl",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-11-14 15:01:26",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "pytorch",
"github_project": "tensordict",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "tensordict"
}
JSON
