How can I build a neural network with 2 independent sets of weights and 2 losses?
Consider the following neural network:
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)

I have 2 independent sets of weights, call them blues (e.g self.blue_1) and reds (e.g self.red_1). These 2 sets need to be multiplied in some combination (e.g see forward method). However, in comparison to what I have above I need the blue weights to be updated according to a certain loss function (e.g loss_blue = some_loss_function, and the red weights to be updated according to another loss function (e.g torch.sum(torch.randint(0,10,(10,)) - x)).
It's important that the red loss doesn't propagate through the blue weights and vice versa.
Is there a way to do this? I was thinking to even break it into 2 neural networks but I'm not sure if 1. it's the best approach, and 2. how to do it using that approach.

One way to do this would be, in four steps:

Do the first inference to compute loss_blue.
Backpropagate from loss_blue and update the blue parameters,
Then infer again to compute loss_red with the updated blue weights
Backpropagate from loss_red and update the red parameters.

This is similar to how you would go about training a GAN alternating between generator and discriminator with successive backward passes and updates. Having two optimizers handling the two sets of parameters makes it easier to use. Don't forget to clear the gradients before backpropagating so one loss doesn't pollute the other parameters with its gradient.

Something like this should work out:
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()


Edit based on comment:


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...])?


Yes, something like that. For defining the optimizers, you can for example create two functions inside your model: reds and blues.
    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()]



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?


When loss_blue.backward, the gradient flows through all parameters of the model, including the red parameters. What makes all the difference is indeed the fact that optim_blue will only update the blue parameters, not the red ones.
