oneflow.nn.Module.state_dict

Module.state_dict(destination=None, prefix='', keep_vars=False)Dict[str, Tensor]

Returns a dictionary containing a whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Parameters
  • destination (dict, optional) – Deprecated. This dict is returned with the module state saved in it. It should also have an attribute _metadata: dict to save metadata of the module state. If it’s not provided, an OrderedDict is created and returned. Default: None

  • prefix (str, optional) – a prefix added to parameter and buffer names to compose the keys in dict. Default: ''

  • keep_vars (bool, optional) – by default the Tensor s returned in the state dict are detached from autograd. If it’s set to True, detaching is not performed. Default: False

Returns

a dictionary containing a whole state of the module

Return type

dict

Example:

>>> import oneflow.nn as nn
>>> l1 = nn.Linear(2, 2)
>>> l2 = nn.Linear(2, 2)
>>> net = nn.Sequential(l1, l2)
>>> net.state_dict().keys()
odict_keys(['0.weight', '0.bias', '1.weight', '1.bias'])