pack_memmap

class tensordict.pack_memmap(prefix: str | pathlib.Path, archive_path: str | pathlib.Path, *, compression: Optional[Union[str, int]] = None, compresslevel: Optional[int] = None)

Packs a memory-mapped tensordict directory into a single-file zip archive.

The archive entries replicate the directory layout, so the operation is the exact inverse of unpack_memmap() (and of a plain unzip). Tensor payloads are stored uncompressed by default and aligned to 64 bytes, which lets load_memmap() memory-map the archive and expose every leaf as a zero-copy view.

Parameters:

• prefix (str or Path) – path to a directory previously produced by memmap() and similar methods.

• archive_path (str or Path) – path of the archive to create.

Keyword Arguments:

• compression (str or int, optional) – one of "stored" (default), "deflate", "bzip2", "lzma" or a zipfile compression constant. Any value other than "stored" trades memory-mapped (zero-copy) loading for a smaller file: compressed leaves are decompressed in memory when accessed.

• compresslevel (int, optional) – compression level forwarded to zipfile.

Returns:

the path to the archive.

Examples

>>> import torch
>>> from tensordict import TensorDict, pack_memmap
>>> td = TensorDict(a=torch.randn(3), b=TensorDict(c=torch.zeros(3, 2)))
>>> td.memmap("./saved_td")
>>> pack_memmap("./saved_td", "./saved_td.tdz")
>>> td2 = TensorDict.load_memmap("./saved_td.tdz")
>>> assert (td == td2).all()