How to seperately backpropogate two loss function

implementation help
1

I am training a model, which has two stages, the stage 1 output is fed to stage 2.out_1 and out_2 are the outputs of stage1 and stage 2 respectively.

I want to backpropagate the losses separately.(My assumption is the loss1 will update the stage 1 and loss 2 will backpropagate and updates stage2 as well as stage 1. bcoz stage 1 output is fed to stage 2)

Please help with this?

code snippet:

def training_step(self, batch, batch_idx):
inputs, targets = batch

output1, output2 = self(inputs)  # Forward pass

# Compute your loss functions
loss1 = sisnr(output1, targets)
loss2 = sisnr(output2, targets)

t_loss = (loss1+loss2)/2

# Log individual losses
self.log('train_loss1', loss1,prog_bar=True)
self.log('train_loss2', loss2,prog_bar=True)
self.log('loss combined',t_loss,prog_bar=True)



# Return a dictionary with individual losses
return {'loss': loss1, 'loss1': loss2}
2

Are you looking to isolate the weight updates of stage1 parameters coming from loss2 vs. from loss1? Or are you interested in preventing backpropagation from loss2 polluting stage1 parameters?

To process the two stages separately you can do something like this:

    output1 = self.stage1(inputs)
    output2 = self.stage2(output1)  

    loss1 = sisnr(output1, targets)  # Loss for stage 1
    loss2 = sisnr(output2, targets)  # Loss for stage 2 (adapting both stages 1 and 2)

    # Zero gradients and create grads from two losses separately but apply them together
    self.optimizer.zero_grad()
    loss1.backward(retain_graph=True) # to keep the updates around
    loss2.backward() 
    self.optimizer.step()

If interested in isolating stage1 from stage2 backpropagation, you may detach out1 before feeding it to stage2:

    output1 = self.stage1(inputs)
    detached_out1 = output1.detach()
    output2 = self.stage2(detached_out1)  

    loss1 = sisnr(output1, targets)  # Loss for stage 1
    loss2 = sisnr(output2, targets)  # Loss for stage 2 (adapting only stage 2)

    # Zero gradients and create grads from two losses separately but apply them together
    self.optimizer.zero_grad()
    loss1.backward(retain_graph=True) # to keep the grads around
    loss2.backward() 
    self.optimizer.step()