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.
pip install tensordictVerified import paths — ran on the pinned version, not inferred.
This quickstart demonstrates the creation of a TensorDict, moving it to a device, slicing it like a tensor, and stacking multiple TensorDicts.
Upgrade Python to version 3.10 or higher.
Replace `td.lock()` with `td.lock_()`, `td.unlock()` with `td.unlock_()`, and `td.rename_key('old', 'new')` with `td.rename_key_('old', 'new')`.Access the tensor content directly, e.g., `my_memmap_tensor` instead of `my_memmap_tensor._tensor`.
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 the library using pip: ```python pip install tensordict ```
To iterate over keys, values, or key-value pairs, explicitly use the `.keys()`, `.values()`, or `.items()` methods:
```python
import tensordict
import torch
td = tensordict.TensorDict({"a": torch.tensor(1), "b": torch.tensor(2)}, batch_size=[])
# Iterate over keys
for key in td.keys():
print(key)
# Iterate over values
for value in td.values():
print(value)
# Iterate over items (key-value pairs)
for key, value in td.items():
print(f"{key}: {value}")
```Ensure that all `TensorDict`s you intend to combine have an identical set of keys. You may need to add missing keys with placeholder values (e.g., zeros) or select a common subset of keys.
```python
import tensordict
import torch
td1 = tensordict.TensorDict({"a": torch.randn(2), "b": torch.randn(2)}, batch_size=[])
td2 = tensordict.TensorDict({"a": torch.randn(2), "c": torch.randn(2)}, batch_size=[])
# This would cause the error: tensordict.stack([td1, td2], dim=0)
# Fix: Ensure keys are identical before stacking
td1_fixed = td1.clone()
td1_fixed["c"] = torch.zeros_like(td2["c"]) # Add missing key to td1
td2_fixed = td2.clone()
td2_fixed["b"] = torch.zeros_like(td1["b"]) # Add missing key to td2
stacked_td = tensordict.stack([td1_fixed, td2_fixed], dim=0)
print(stacked_td)
```Before performing such operations, explicitly move all relevant `TensorDict`s or their contained tensors to the target device using the `.to()` method.
```python
import tensordict
import torch
td_cpu = tensordict.TensorDict({"a": torch.randn(2)}, batch_size=[])
if torch.cuda.is_available():
device = "cuda:0"
td_cuda = tensordict.TensorDict({"a": torch.randn(2).to(device)}, batch_size=[])
# This would cause the RuntimeError if uncommented:
# combined_td = tensordict.stack([td_cpu, td_cuda], dim=0)
# Fix: Move one or both TensorDicts to the same device
td_cpu_on_cuda = td_cpu.to(device) # Move td_cpu to GPU
combined_td = tensordict.stack([td_cpu_on_cuda, td_cuda], dim=0)
print(combined_td)
else:
print("CUDA not available, skipping device specific example.")
```Apply the operation to specific tensors within the `TensorDict` (e.g., `td["key"] + scalar`) or use `tensordict.apply()` for element-wise modification of all tensors.