我的神经网络有以下结构:
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层,而我希望这个操作可以并行完成。