如何构建具有 2 个独立权重集和 2 个损失的神经网络?

How can I build a neural network with 2 independent sets of weights and 2 losses?

考虑以下神经网络:

import torch
import torch.nn as nn
import torch.optim as optim

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.blue_1 = nn.Linear(2, 10)
        self.red_1 = nn.Linear(10, 5)
        self.blue_2 = nn.Linear(5, 4)
        self.red_2 = nn.Linear(4, 3)
        self.blue_3 = nn.Linear(3, 2)
        self.red_3 = nn.Linear(2, 1)

    def forward(self, x):
        x = torch.relu(self.blue_1(x))
        x = self.red_1(x)
        x = self.blue_2(x)
        x = self.red_2(x)
        x = self.blue_3(x)
        x = self.red_3(x)
        return x

net = Model()

opt = optim.Adam(net.parameters())

features = torch.rand((10,2)) #10 inputs, each of 2D

for epoch in range(3):
    x = net(features)
    loss = torch.sum(torch.randint(0,10,(10,)) - x)
    loss.backward()
    print(loss)

我有 2 组独立的权重,称它们为 blues(例如 self.blue_1)和 reds(例如 self.red_1)。这 2 个集合需要以某种组合相乘(例如,参见 forward 方法)。但是,与我上面的相比,我需要根据特定损失函数更新 blue 权重(例如 loss_blue = some_loss_function,并根据另一个损失更新 red 权重功能(例如 torch.sum(torch.randint(0,10,(10,)) - x))。

重要的是 red 损失不会通过 blue 权重传播,反之亦然。

有办法吗?我什至想把它分成 2 个神经网络,但我不确定 1. 它是否是最好的方法,以及 2. 如何使用该方法来实现。

一种方法是分四步完成:

  1. 进行第一个推理以计算 loss_blue
  2. loss_blue反向传播并更新蓝色参数,
  3. 然后再次推断以使用更新后的蓝色权重
  4. 计算loss_red
  5. loss_red 反向传播并更新红色参数。

这类似于您通过连续向后传递和更新在生成器和鉴别器之间交替训练 GAN 的方式。让两个优化器处理两组参数使其更易于使用。不要忘记在反向传播之前清除梯度,这样一个损失就不会用它的梯度污染其他参数。


这样的事情应该可以解决:

net = Model()

optim_blue = optim.Adam(net.blues()) # fetch all blue parameters
optim_red = optim.Adam(net.reds())   # fetch all red parameters

features = torch.rand((10,2)) #10 inputs, each of 2D

# # inference, backprop and update on blue params
out = net(features)
loss_blue = torch.sum(torch.randint(0,10,(10,)) - out)

optim_red.zero_grad()
optim_blue.zero_grad()
loss_blue.backward()
optim_blue.step()

# # inference backprop and update on red params
out = net(features)
loss_red = out.mean()

optim_red.zero_grad()
optim_blue.zero_grad()
loss_red.backward()
optim_red.step()

根据评论编辑:

  1. How do you specify which parameters/layers will be in the optimizer(optim_blue = optim.Adam(net.blues())). Is it something along the lines of optim_blue = optim.Adam([net. blue_1, net. blue_2...])?

是的,类似的东西。为了定义优化器,您可以在模型中创建两个函数:redsblues.

    def reds(self):
        return [*self.red_1.parameters(),
                *self.red_2.parameters(),
                *self.red_3.parameters()]

    def blues(self):
        return [*self.blue_1.parameters(),
                *self.blue_2.parameters(),
                *self.blue_3.parameters()]
  1. What makes the gradients separated when you call loss_blue.backward()? That is, what stops them to flow through the reds? Is that the purpose of the 2 optimizers?

loss_blue.backward时,梯度流过模型的所有参数,包括红色参数。真正不同的是 optim_blue 只会更新蓝色参数,不会更新红色参数。