WrapModule

class tensordict.nn.WrapModule(*args, **kwargs)

一个围绕任何处理 TensorDict 实例的可调用对象的包装器。

当构建 TensorDictSequential 堆栈且某个转换需要整个 TensorDict 实例可见时，此包装器很有用。

参数:

func (Callable[[TensorDictBase], TensorDictBase]) – 一个可调用函数，它接受一个 TensorDictBase 实例并返回一个转换后的 TensorDictBase 实例。

关键字参数:

• inplace (bool, optional) – 如果为 True，则输入 TensorDict 将被原地修改。否则，将返回一个新的 TensorDict（如果函数不原地修改它并返回它）。默认为 False。

• in_keys (list of NestedKey, optional) – 如果提供，则指示模块读取哪些条目。这不会被检查，仅用于告知 TensorDictSequential 被包装模块的输入键。默认为 []。

• out_keys (list of NestedKey, optional) – 如果提供，则指示模块写入哪些条目。这不会被检查，仅用于告知 TensorDictSequential 被包装模块的输出键。默认为 []。

示例

>>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod, WrapModule
>>> seq = Seq(
...     Mod(lambda x: x * 2, in_keys=["x"], out_keys=["y"]),
...     WrapModule(lambda td: td.reshape(-1)),
... )
>>> td = TensorDict(x=torch.ones(3, 4, 5), batch_size=[3, 4])
>>> td = Seq(td)
>>> assert td.shape == (12,)
>>> assert (td["y"] == 2).all()
>>> assert td["y"].shape == (12, 5)

forward(data: TensorDictBase) → TensorDictBase

定义每次调用时执行的计算。

所有子类都应重写此方法。

注意

尽管前向传播的实现需要在此函数中定义，但您应该在之后调用 Module 实例而不是此函数，因为前者会处理注册的钩子，而后者则会静默忽略它们。