reference/generated/tensordict.append_memmap

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

append_memmap

tensordict.append_memmap(prefix: str | Path, rows: TensorCollection, *, allow_pickle: bool | None = None) → None

Appends rows to a tensordict saved in a directory, along its first batch dimension.

Each leaf file grows by the bytes of the new rows, and the meta.json files are rewritten last, each one through a temporary file, so that a failed append leaves the saved data as it was. Readers that mapped the files before the append keep seeing the old number of rows, and load_memmap() after the append sees the new one. Use one writer at a time. The meta.json files of nested levels are replaced one at a time, so an append that is killed midway, or a load_memmap() that runs during an append, can see nested levels with different numbers of rows: don’t read a directory while you append to it. A killed append can leave its new meta.json files behind as .meta.json.append.tmp files. They are all written before the first meta.json is replaced, and the top meta.json, in the _tensordict subdirectory for a tensorclass, is replaced first. If .meta.json.append.tmp is still in the top directory, no meta.json was replaced: delete every temporary file, which may be incomplete, and the directory holds the old rows. Otherwise, rename each temporary file over the meta.json next to it, which completes the append. Decide by this file, not by whether load_memmap() raises, which it doesn’t for every tree. Each append maps the whole saved tree, as load_memmap() does, to read its structure, so many small appends to a save with many leaves are slow. See Save and load tensordicts.

rows must have the keys, the classes, the dtypes and the trailing shapes of the saved tensordict, and its batch size may differ from the saved one in its first dimension only. A non-tensor entry stays as saved, and its value in rows must be equal to it: either one value or a NonTensorStack of equal values. A saved NonTensorStack, a lazy stack, a nested tensor and an archive cannot grow, and raise. Everything is checked before the first file is written.

Parameters:

• prefix (str or Path) – the directory that save() or memmap() wrote.

• rows (TensorDictBase or tensorclass) – the rows to append.

Keyword Arguments:

allow_pickle (bool, optional) – whether the pickled non-tensor entries of the saved tensordict, which are compared with those of rows, may be unpickled, as in load_memmap(). Defaults to None.

Raises:

ValueError – if prefix is not a directory, or if rows cannot be appended to the saved tensordict.

Examples

>>> import tempfile
>>> import torch
>>> import tensordict
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3, 2), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as tmpdir:
...     _ = tensordict.save(td, tmpdir)
...     tensordict.append_memmap(tmpdir, TensorDict(a=torch.ones(2, 2), batch_size=[2]))
...     print(tensordict.load_memmap(tmpdir)["a"].sum(-1).tolist())
[0.0, 0.0, 0.0, 2.0, 2.0]