reference/generated/tensordict.TensorDictBase.state_dict

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.state_dict

TensorDictBase.state_dict(destination=None, prefix='', keep_vars=False, flatten=True) → OrderedDict[str, Any]

Returns the leaves of the tensordict in an OrderedDict, as torch.nn.Module.state_dict() does.

The state dict is flat by default: the key of a nested leaf joins its keys with "." ("n.b" for ("n", "b")), following torch.nn.Module.state_dict(). With flatten=False, each nested tensordict becomes a nested OrderedDict. Read it back with load_state_dict().

The batch size and device of each level are stored in a _metadata attribute of the result, as in torch.nn.Module.state_dict(). In a flat state dict, they are stored under the "."-joined key of the level ("" for the root, "n" for the nested tensordict "n"). In a nested state dict, each nested OrderedDict has its own _metadata.

Parameters:

• destination (OrderedDict, optional) – an OrderedDict to write the entries into, and to return. A plain dict raises an AttributeError. Defaults to None, which creates a new OrderedDict.

• prefix (str, optional) – a string prepended to every key. Defaults to "".

• keep_vars (bool, optional) – if False, the tensors are detached from autograd. If True, the tensors of the tensordict are returned unchanged. Defaults to False.

• flatten (bool, optional) – if True, nested keys are joined with ".". If False, nested tensordicts become nested OrderedDict instances. Defaults to True.

Raises:

KeyError – if flatten=True and two entries get the same key, as a dotted key "a.b" and a nested key ("a", "b") do. The key of a level in _metadata counts too. Use flatten=False to keep such entries apart.

Returns:

destination, or a new OrderedDict.

Examples

>>> from tensordict import TensorDict
>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, [])
>>> sd = data.state_dict()
>>> print(dict(sd))
{'1': tensor(1), '2': tensor(2), '3.3': tensor(3)}
>>> print(dict(sd._metadata))
{'': {'batch_size': torch.Size([]), 'device': None}, '3': {'batch_size': torch.Size([]), 'device': None}}

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