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)