TensorDict

0.12.0 · active · verified Sat Apr 11

TensorDict is a PyTorch-dedicated tensor container, offering a dictionary-like class that inherits properties from `torch.Tensor`. It simplifies working with collections of tensors, enabling tensor-like operations, efficient data management, and re-usable training loops across various machine learning paradigms. It's currently at version 0.12.0 and is actively developed.

Warnings

breaking Python 3.9 support was dropped in TensorDict v0.11.0. Python 3.10 or newer is now required.
Fix: Upgrade Python to version 3.10 or higher.
breaking Deprecated methods `lock`, `unlock`, and `rename_key` (without a trailing underscore) were removed in v0.11.0. Use `lock_`, `unlock_`, and `rename_key_` instead for in-place modifications.
Fix: Replace `td.lock()` with `td.lock_()`, `td.unlock()` with `td.unlock_()`, and `td.rename_key('old', 'new')` with `td.rename_key_('old', 'new')`.
breaking The `MemoryMappedTensor._tensor` property now raises a `RuntimeError` since v0.11.0. Users should interact with the `MemoryMappedTensor` instance directly as it is a tensor subclass.
Fix: Access the tensor content directly, e.g., `my_memmap_tensor` instead of `my_memmap_tensor._tensor`.
gotcha From v0.10.0, lists assigned to a TensorDict will be automatically stacked by default, potentially raising a `FutureWarning`. Explicit context managers should be used for specific behavior.
Fix: Review code that assigns lists to TensorDicts and explicitly use context managers like `td.set_` if list stacking is not the desired default behavior, or if you want to suppress the warning.

Install

pip install tensordict PyPI
conda install -c conda-forge tensordict Conda

Imports

TensorDict
```
from tensordict import TensorDict
```
tensorclass
```
from tensordict import tensorclass
```

MemoryMappedTensor

from tensordict import MemoryMappedTensor

Quickstart

This quickstart demonstrates the creation of a TensorDict, moving it to a device, slicing it like a tensor, and stacking multiple TensorDicts.

import torch
from tensordict import TensorDict

data = TensorDict(
    obs=torch.randn(128, 84),
    action=torch.randn(128, 4),
    reward=torch.randn(128, 1),
    batch_size=[128],
)

device = "cuda" if torch.cuda.is_available() else "cpu"
data_gpu = data.to(device) # all tensors move together
sub = data_gpu[:64] # all tensors are sliced
stacked = torch.stack([data, data]) # works like a tensor

print(f"Original TensorDict:\n{data}")
print(f"Device of data_gpu: {data_gpu.device}")
print(f"Batch size of sliced TensorDict: {sub.batch_size}")
print(f"Batch size of stacked TensorDict: {stacked.batch_size}")

view raw JSON →