from_module

class tensordict.from_module(module, as_module: bool = False, lock: bool = True, use_state_dict: bool = False)

Copies the params and buffers of a module in a tensordict.

Parameters:

• module (nn.Module) – the module to get the parameters from.

• as_module (bool, optional) – if True, a TensorDictParams instance will be returned which can be used to store parameters within a torch.nn.Module. Defaults to False.

• lock (bool, optional) – if True, the resulting tensordict will be locked. Defaults to True.

• use_state_dict (bool, optional) –

  if True, the state-dict from the module will be used and unflattened into a TensorDict with the tree structure of the model. Defaults to False.

  Note

  This is particularly useful when state-dict hooks have to be used.

Examples

>>> from torch import nn
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1)
>>> params = from_module(module)
>>> print(params["layers", "0", "linear1"])
TensorDict(
    fields={
        bias: Parameter(shape=torch.Size([2048]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Parameter(shape=torch.Size([2048, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)