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

Question

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

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

我使用一个ModuleList来保存完全连接层的列表。以下是对这一点的看法：

^{pr2}$

然后，当我需要计算输出时，我使用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参数来表示“分组卷积”（又称“depthwise卷积”）。让你同时处理所有的网络。在

如果使用一个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)

以及

^{pr2}$

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

---

虽然CUDA本机支持分组卷积，但是Pythorch在分组卷积的速度方面存在一些问题（参见示例here），但我认为现在已经解决了。在