与 PyTorch 并行运行一个集成的多个模型

Question

我的神经网络具有以下架构：

input -> 128x (separate fully connected layers) -> output averaging

我正在使用 ModuleList 来保存完全连接层的列表。这是它在这一点上的样子：

class MultiHead(nn.Module):
    def __init__(self, dim_state, dim_action, hidden_size=32, nb_heads=1):
        super(MultiHead, self).__init__()

        self.networks = nn.ModuleList()
        for _ in range(nb_heads):
            network = nn.Sequential(
                nn.Linear(dim_state, hidden_size),
                nn.Tanh(),
                nn.Linear(hidden_size, dim_action)
            )
            self.networks.append(network)

        self.cuda()
        self.optimizer = optim.Adam(self.parameters())

然后，当我需要计算输出时，我使用一个for ... in构造来执行所有层的前向和后向传递：

q_values = torch.cat([net(observations) for net in self.networks])

# skipped code which ultimately computes the loss I need

self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()

这有效！但我想知道我是否不能更有效地做到这一点。我觉得通过做一个for...in，我实际上是在逐个浏览每个单独的 FC 层，而我希望这个操作可以并行完成。

Answer 1

在Convnd代替的情况下，Linear您可以使用groups“分组卷积”（又名“深度卷积”）的参数。这让您可以同时处理所有并行网络。

如果您使用大小为 的卷积核1，则卷积除了应用一个Linear层之外什么都不做，其中每个通道都被视为一个输入维度。因此，您的网络的粗略结构如下所示：

1. 修改形状的输入张量B x dim_state如下：添加一个额外的维度并复制nb_state-times B x dim_state到B x (dim_state * nb_heads) x 1
2. 代替两个Linear用

nn.Conv1d(in_channels=dim_state * nb_heads, out_channels=hidden_size * nb_heads, kernel_size=1, groups=nb_heads)

和

nn.Conv1d(in_channels=hidden_size * nb_heads, out_channels=dim_action * nb_heads, kernel_size=1, groups=nb_heads)

4. 我们现在有一个大小的张量，B x (dim_action x nb_heads) x 1您现在可以将其修改为您想要的任何形状（例如B x nb_heads x dim_action）

虽然 CUDA 本身支持分组卷积，但 pytorch 中分组卷积的速度存在一些问题（参见例如此处），但我认为现在已经解决了。