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层外无所作为，在该层中，每个 channel 被视为一个输入维。因此，您的网络的粗略结构如下所示：

按如下所示修改形状为B x dim_state的输入张量：添加一个附加维度，并以nb_state倍B x dim_state的形式复制到B x (dim_state * nb_heads) x 1
将两个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中存在一些问题，例如分组卷积的速度（请参见here），但我认为现在已经解决了。

与PyTorch并行运行集合的多个模型

1 个答案: