使用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层，而我希望这个操作可以并行完成。

- MasterScrat

Answer 1

在使用“组卷积”（也称为“深度卷积”）时，可以将Convnd用于Linear，并使用groups参数来处理所有并行网络。如果使用大小为1的卷积核，则卷积除了应用每个通道作为输入维度的Linear层外，没有其他操作。因此，您的网络的大致结构如下所示：

1. 将形状为B x dim_state的输入张量修改如下：添加一个额外的维度，并通过nb_state次复制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)

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

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

- flawr

---

听起来很有前途，我会试一下！不过我认为你的第一个代码片段可能存在复制粘贴错误;-) - MasterScrat

哎呀，我不知道那个是怎么进去的，抱歉！ - flawr