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. The next append_memmap to the directory finds these files first. If .meta.json.append.tmp is gone from the top directory and the top meta.json already has the number of rows of every temporary file, it renames each one over the meta.json next to it, which completes the killed append. Otherwise no meta.json was replaced, or the rollback of a failed append put the old ones back: it deletes every temporary file, which may be incomplete, and the directory holds the old rows. To read the directory without appending, apply the same rule yourself (see Save and load tensordicts), and decide by these files, 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]