- 问题:
-
我的神经网络具有以下结构：
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())然后，当我需要计算输出时，我用一个
来表示。。。在构造中执行所有层的正向和反向传递：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()这很管用！但我想知道我是否能更有效地完成这项工作。我觉得通过在中对…执行
操作，我实际上是在逐个遍历每个单独的FC层，而我希望这个操作可以并行完成
- 答案:
-
在
convd代替线性的情况下，可以使用groups参数来表示“分组卷积”（又称“depthwise卷积”）。这让你同时处理所有的并行网络如果使用大小为
1的卷积内核，那么卷积只会应用线性层，其中每个通道都被视为输入维度。所以你的网络结构大致如下：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)虽然CUDA本机支持分组卷积，但Pythorch在分组卷积速度方面存在一些问题（参见。here)但我想现在已经解决了
 |
 |
 |