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 可迭代对象给出的顺序传递。
out_keys (iterable of str) – 要写入输入 tensordict 的键。out_keys 的长度必须与嵌入式模块返回的张量数量匹配。使用“_”作为键可以避免将张量写入输出。
spec (TensorSpec, optional) – 输出张量的规格。如果模块输出多个输出张量，spec 则表征第一个输出张量的空间。
safe (bool) – 如果为 True，则输出的值会根据输入 spec 进行检查。由于探索策略或数值下溢/溢出问题，可能会发生域外采样。如果此值超出边界，将使用 TensorSpec.project 方法将其投影回所需空间。默认为 False。
inplace (bool or str, optional) – 如果为 True，则输入 tensordict 会被就地修改。如果为 False，则会创建一个新的空 TensorDict 实例。如果为 “empty”，则使用 input.empty()（即输出保留类型、设备和批次大小）。默认为 True。

在 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: torch.dtype | DEVICE_TYPING) → 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 张量移动到 CUDA 设备。

有关示例，请参阅下文。

注意

此方法就地修改模块。

参数:

device (torch.device) – 此模块中的参数和缓冲区的期望设备
dtype (torch.dtype) – 此模块中的参数和缓冲区的期望浮点或复数 dtype
tensor (torch.Tensor) – 具有此模块所有参数和缓冲区的期望 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¶

文档

教程

资源