memmap#

tensordict.memmap(data: TensorCollection, prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True, robust_key: bool | None = True, archive: bool | None = None, compression: str | int | None = None, atomic: bool = False) → TensorCollection#

Writes every leaf to a memory-mapped file and returns a new tensordict of them.

Same as save(), and as memmap(), which holds the full text of every argument. Use it to share a tensordict between processes without copies, or to keep data that is larger than RAM on disk. To convert the tensordict in place, use memmap_(). See Save and load tensordicts and Load data and share it between processes.

Parameters:

data (TensorDictBase or tensorclass instance) – the tensordict to write.
prefix (str, optional) – the directory to write to, or a file name that ends with ".tdz". Defaults to None, which writes the tensors to anonymous shared memory and saves nothing to disk.
copy_existing (bool, optional) – if False, an entry that is already memory-mapped to a file outside prefix raises an error. If True, it is copied. Defaults to False.

Keyword Arguments:

num_threads (int, optional) – the number of threads that write the files. Defaults to 0.
return_early (bool, optional) – if True and num_threads > 1, returns a TensorDictFuture at once, while the threads write the files. Defaults to False.
share_non_tensor (bool, optional) – if True, non-tensor entries are shared between the processes of a node. Defaults to False.
existsok (bool, optional) – if False, a tensor file that already exists under prefix raises an error. Defaults to True.
robust_key (bool, optional) – if True, keys that hold path separators or special characters are encoded in the file names. Defaults to True.
archive (bool, optional) – if True, prefix is a file and the tensordict is written as a single zip archive that load() can memory-map. If None, an archive is written when prefix ends with ".tdz". Defaults to None.
compression (str or int, optional) – the compression of the archive entries: "stored", "deflate", "bzip2", "lzma" or a zipfile constant. Only valid for an archive. Defaults to None, which stores them uncompressed. Only uncompressed entries load without a copy.
atomic (bool, optional) – if True, the directory is written next to prefix and then renamed to it, so that an error or a crash during the write leaves prefix as it was. Not supported on Windows for a directory. memmap() gives the details. Defaults to False.

Returns:

A new, locked tensordict whose leaves are MemoryMappedTensor instances.

Examples

>>> import torch
>>> import tensordict
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3), batch_size=[3])
>>> td_memmap = tensordict.memmap(td)
>>> type(td_memmap["a"]).__name__
'MemoryMappedTensor'
>>> td_memmap.is_locked
True

memmap#

Docs

Tutorials

Resources