SafeModule¶

class torchrl.modules.tensordict_module.SafeModule(*args, **kwargs)[原始碼]¶

tensordict.nn.TensorDictModule 子類別，接受 TensorSpec 作為參數以控制輸出域。

參數:

module (nn.Module) – 用於將輸入映射到輸出參數空間的 nn.Module。可以是功能模組 (FunctionalModule 或 FunctionalModuleWithBuffers)，在這種情況下，forward 方法將期望 params（和可能的）buffers 關鍵字引數。
in_keys (iterable of str) – 要從輸入 tensordict 讀取並傳遞到模組的鍵。如果它包含多個元素，則這些值將按照 in_keys iterable 給出的順序傳遞。
out_keys (iterable of str) – 要寫入輸入 tensordict 的鍵。out_keys 的長度必須與嵌入式模組傳回的張量數量相符。使用 "_" 作為鍵可避免將張量寫入輸出。
spec (TensorSpec, optional) – 輸出張量的規格。如果模組輸出多個輸出張量，則 spec 會描述第一個輸出張量的空間。
safe (bool) – 如果 True，則根據輸入規格檢查輸出的值。由於探索策略或數值下溢/溢位問題，可能會發生超出範圍的取樣。如果此值超出範圍，則使用 TensorSpec.project 方法將其投影回所需的空間。預設值為 False。

在 TensorDictModule 中嵌入神經網路只需要指定輸入和輸出鍵。如果需要，可以傳遞域規格。: TensorDictModule 支援功能性和常規 nn.Module 物件。在功能性情況下，必須指定 'params'（和 'buffers'）關鍵字引數

範例

>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.data import Unbounded
>>> from torchrl.modules import TensorDictModule
>>> td = TensorDict({"input": torch.randn(3, 4), "hidden": torch.randn(3, 8)}, [3,])
>>> spec = Unbounded(8)
>>> module = torch.nn.GRUCell(4, 8)
>>> td_fmodule = TensorDictModule(
...    module=module,
...    spec=spec,
...    in_keys=["input", "hidden"],
...    out_keys=["output"],
...    )
>>> params = TensorDict.from_module(td_fmodule)
>>> with params.to_module(td_module):
...     td_functional = td_fmodule(td.clone())
>>> print(td_functional)
TensorDict(
    fields={
        hidden: Tensor(torch.Size([3, 8]), dtype=torch.float32),
        input: Tensor(torch.Size([3, 4]), dtype=torch.float32),
        output: Tensor(torch.Size([3, 8]), dtype=torch.float32)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

在具狀態的情況下

>>> td_module = TensorDictModule(
...    module=torch.nn.GRUCell(4, 8),
...    spec=spec,
...    in_keys=["input", "hidden"],
...    out_keys=["output"],
...    )
>>> td_stateful = td_module(td.clone())
>>> print(td_stateful)
TensorDict(
    fields={
        hidden: Tensor(torch.Size([3, 8]), dtype=torch.float32),
        input: Tensor(torch.Size([3, 4]), dtype=torch.float32),
        output: Tensor(torch.Size([3, 8]), dtype=torch.float32)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

可以使用 vmap 運算子來呼叫功能模組。在這種情況下，tensordict 會擴充以符合批次大小（即，tensordict 不再就地修改）

>>> # Model ensemble using vmap
>>> from torch import vmap
>>> params_repeat = params.expand(4, *params.shape)
>>> td_vmap = vmap(td_fmodule, (None, 0))(td.clone(), params_repeat)
>>> print(td_vmap)
TensorDict(
    fields={
        hidden: Tensor(torch.Size([4, 3, 8]), dtype=torch.float32),
        input: Tensor(torch.Size([4, 3, 4]), dtype=torch.float32),
        output: Tensor(torch.Size([4, 3, 8]), dtype=torch.float32)},
    batch_size=torch.Size([4, 3]),
    device=None,
    is_shared=False)

random(tensordict: TensorDictBase) → TensorDictBase[原始碼]¶

在目標空間中取樣一個隨機元素，無論任何輸入如何。

如果存在多個輸出鍵，則只有第一個會寫入輸入 tensordict。

參數:: tensordict (TensorDictBase) – 應該寫入輸出值的 tensordict。
回傳值:: 帶有輸出鍵的新/更新值的原始 tensordict。

random_sample(tensordict: TensorDictBase) → TensorDictBase[原始碼]¶: 請參閱 TensorDictModule.random(...)。

to(dest: Union[dtype, device, str, int]) → TensorDictModule[原始碼]¶

移動和/或轉換參數和緩衝區。

可以這樣呼叫：

to(device=None, dtype=None, non_blocking=False)[原始碼]

to(dtype, non_blocking=False)[原始碼]

to(tensor, non_blocking=False)[原始碼]

to(memory_format=torch.channels_last)[原始碼]

其簽名檔與 torch.Tensor.to() 相似，但僅接受浮點數或複數 dtype。此外，此方法僅會將浮點數或複數參數和緩衝區轉換為 dtype（如果已給定）。整數參數和緩衝區將會被移動到 device（如果已給定），但 dtype 不會變更。當設定 non_blocking 時，如果可能，它會嘗試相對於主機非同步轉換/移動，例如，將具有固定記憶體的 CPU Tensor 移動到 CUDA 裝置。

請參閱以下範例。

注意

此方法會就地修改模組。

參數:

device (torch.device) – 此模組中參數和緩衝區的所需裝置
dtype (torch.dtype) – 此模組中參數和緩衝區的所需浮點數或複數 dtype
tensor (torch.Tensor) – Tensor，其 dtype 和裝置是此模組中所有參數和緩衝區的所需 dtype 和裝置
memory_format (torch.memory_format) – 此模組中 4D 參數和緩衝區的所需記憶體格式（僅限關鍵字引數）

回傳值:

self

回傳類型:

模組

範例

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

SafeModule¶

文件

教學課程

資源