reference/generated/tensordict.load

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

load

tensordict.load(prefix: str | Path, device: device | None = None, non_blocking: bool = False, *, out: TensorCollection | None = None, robust_key: bool | None = True, subpath: NestedKey | None = None, mode: str | None = None, num_threads: int = 0, allow_pickle: bool | None = None) → TensorCollection

Loads a tensordict that save() wrote.

Same as load_memmap(), which takes the same arguments. See tensordict.TensorDict.load_memmap() for every argument, and Save and load tensordicts for the ways to save and load a tensordict.

Parameters:

• prefix (str or Path) – the directory, or the .tdz archive file, that holds the saved tensordict.

• device (device, str or int, optional) – if given, the data is cast to that device. With "meta", no file is read. Defaults to None.

• non_blocking (bool, optional) – if True, the cast to device is not followed by a synchronization. Defaults to False.

Keyword Arguments:

• out (TensorDictBase, optional) – a tensordict to load the data into. Defaults to None.

• robust_key (bool, optional) – if True, reads the percent-encoded file names that save() writes by default, and falls back to the legacy names of keys that are safe file names. If False, reads the legacy layout only. Defaults to True.

• subpath (NestedKey or str, optional) – the nested tensordict to load, as a nested key (("module", "0")) or a "/"-separated path ("module/0"). Only that subtree is read. Defaults to None, which loads the whole tensordict.

• mode (str, optional) – "r" maps the files copy-on-write (in-place writes stay in memory), "r+" maps them shared (in-place writes reach the files). See tensordict.TensorDict.load_memmap() for the default.

• num_threads (int, optional) – the number of threads that decompress the leaves of a compressed archive. It has no effect on other data. Defaults to 0.

• allow_pickle (bool, optional) – if True, non-tensor entries that were pickled are unpickled. Only load pickles from trusted sources. If False, pickled entries raise a RuntimeError. Defaults to None, which unpickles them with a FutureWarning.

Returns:

The saved tensordict, of the class it had when it was saved. Its leaves are MemoryMappedTensor instances.

Examples

>>> import tempfile
>>> import torch
>>> import tensordict
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3), b={"c": torch.ones(3, 2)}, batch_size=[3])
>>> with tempfile.TemporaryDirectory() as tmpdir:
...     _ = tensordict.save(td, tmpdir)
...     loaded = tensordict.load(tmpdir)
...     print(loaded["b", "c"].shape)
torch.Size([3, 2])