tensordict.nn.EnsembleModule

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

封装模块并重复它以形成一个集成体的模块。

参数：

• module (nn.Module) – 要复制和封装的 nn.module。

• num_copies (int) – 要创建的 module 副本数量。

• parameter_init_function (Callable) – 一个接收模块副本并初始化其参数的函数。

• expand_input (bool) – 是否将输入 TensorDict 扩展以匹配副本数量。 除非将多个集成模块链接在一起，例如 EnsembleModule(cnn) -> EnsembleModule(mlp)，否则应为 True。 如果为 False，EnsembleModule(mlp) 会期望之前的模块已扩展输入。

示例

>>> import torch
>>> from torch import nn
>>> from tensordict.nn import TensorDictModule, EnsembleModule
>>> from tensordict import TensorDict
>>> net = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 2))
>>> mod = TensorDictModule(net, in_keys=['a'], out_keys=['b'])
>>> ensemble = EnsembleModule(mod, num_copies=3)
>>> data = TensorDict({'a': torch.randn(10, 4)}, batch_size=[10])
>>> ensemble(data)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([3, 10, 2]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3, 10]),
    device=None,
    is_shared=False)

为了将多个 EnsembleModules 堆叠在一起，我们应该注意从第二个模块开始关闭 `expand_input` 并开启。

示例

>>> import torch
>>> from tensordict.nn import TensorDictModule, TensorDictSequential, EnsembleModule
>>> from tensordict import TensorDict
>>> module = TensorDictModule(torch.nn.Linear(2,3), in_keys=['bork'], out_keys=['dork'])
>>> next_module = TensorDictModule(torch.nn.Linear(3,1), in_keys=['dork'], out_keys=['spork'])
>>> e0 = EnsembleModule(module, num_copies=4, expand_input=True)
>>> e1 = EnsembleModule(next_module, num_copies=4, expand_input=False)
>>> seq = TensorDictSequential(e0, e1)
>>> data = TensorDict({'bork': torch.randn(5,2)}, batch_size=[5])
>>> seq(data)
TensorDict(
    fields={
        bork: Tensor(shape=torch.Size([4, 5, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        dork: Tensor(shape=torch.Size([4, 5, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        spork: Tensor(shape=torch.Size([4, 5, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([4, 5]),
    device=None,
    is_shared=False)