使用 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 中存在分组卷积速度方面的一些问题（参见例如here），但我认为现在已经解决了。