from_module

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

将模块的参数和缓冲区复制到 tensordict 中。

参数:

• module (nn.Module) – 从中获取参数的模块。

• as_module (bool, optional) – 如果为 True，将返回一个 TensorDictParams 实例，该实例可用于在 torch.nn.Module 中存储参数。默认为 False。

• lock (bool, optional) – 如果为 True，则生成的 tensordict 将被锁定。默认为 True。

• use_state_dict (bool, optional) –

  如果为 True，则将使用模块的 state-dict 并将其解展平为具有模型树结构的 TensorDict。默认为 False。

  注意

  当必须使用 state-dict 钩子时，这尤其有用。

示例

>>> 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)