TensorDictBase.from_schema#

classmethod TensorDictBase.from_schema(schema: dict[str, tuple[list[int] | Size, dtype]], *, batch_size: Sequence[int] | Size | None = None, storage: str | None = None, device=None, **kwargs) → TensorDictBase#

Allocates a zero-filled tensordict from a schema of shapes and dtypes.

Each entry of schema maps a key to an (element_shape, dtype) pair, and the entry gets the shape [*batch_size, *element_shape]. storage selects where the zeros are stored, so that a buffer on disk or on a server can be created without building it in memory first.

Parameters:

schema (dict) – maps each key (a str) to a pair (element_shape, dtype), where element_shape (a list of int or a torch.Size) is the shape of the entry after the batch dimensions.

Keyword Arguments:

batch_size (Size or list of int, optional) – the batch size of the result, prepended to every element shape. Defaults to None, an empty batch size.
storage (str, optional) –
where the entries are stored:
- None: a TensorDict of regular tensors.
- "memmap": memory-mapped tensors, in the directory given by prefix=<dir> (see memmap_like()).
- "h5": a PersistentTensorDict in the HDF5 file given by filename=<path>.
- "zarr": a PersistentTensorDict in the zarr store given by filename=<path or store> (requires zarr>=3.1.1).
- "shared": tensors in shared memory (see share_memory_()).
- "redis" or "dragonfly": a TensorDictStore on a server, through TensorDictStore.from_schema.
Defaults to None.
device (device, str or int, optional) – the device of the result. The "memmap" storage ignores it. Defaults to None.
**kwargs – passed to the backend: prefix and the other arguments of memmap_like() for "memmap", filename for "h5" and "zarr", host, port and the like for "redis" and "dragonfly".

Returns:

A new tensordict of the type that storage selects, filled with zeros.

Raises:

ValueError – if storage is not one of the values above.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict.from_schema(
...     {"obs": ([84, 84, 3], torch.uint8),
...      "reward": ([], torch.float32)},
...     batch_size=[1000],
... )
>>> td["obs"].shape
torch.Size([1000, 84, 84, 3])

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as d:
...     td_mm = TensorDict.from_schema(
...         {"obs": ([4], torch.float32)},
...         batch_size=[8],
...         storage="memmap",
...         prefix=d,
...     )
...     print(td_mm.is_memmap())
True