reference/generated/tensordict.TensorDictBase.save

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.save

TensorDictBase.save(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, robust_key: bool | None = True, archive: bool | None = None, compression: str | int | None = None) → Self

Saves the tensordict to disk, and returns a memory-mapped copy of it.

save (and its alias dumps) is memmap(): it takes the same arguments. self is not changed. Read the result back with load(). See Save and load tensordicts.

Parameters:

prefix (str or Path) – the directory where the tensordict is written, one file per leaf. If prefix ends with ".tdz" or archive=True, a single archive file is written instead.

Keyword Arguments:

• archive (bool, optional) – if True, write a single archive file. Defaults to None, which writes an archive if prefix ends with ".tdz".

• compression (str or int, optional) – the compression of the archive. See memmap().

All the other arguments are those of memmap().

Examples

>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3), nested=TensorDict(b=torch.ones(3), batch_size=[3]), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as path:
...     _ = td.save(path)
...     loaded = TensorDict.load(path)
...     print((loaded == td).all())
True

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

Tutorials that use save

Benchmarking TensorDict serialization speed

Benchmarking TensorDict serialization speed