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)