[how-to] Handle multiple losses and/or weighted losses?
What is your question?
I'm trying replicate the model built in https://github.com/ekagra-ranjan/AE-CNN with Pytorch Lightning as a way to learn the framework.
They use an Auto-encoder along with a CNN (e.g. Inception V3), and this means there are multiple loss functions for each model, with a separate weight for each:
loss1 = MSE loss for auto-encoder
loss2 = 0.8 * BCE loss for one branch of inceptionv3 + 0.2 * BCE loss for another branch of inceptionv3
overall loss = 0.1 * loss1 + 0.9 * loss2

How would I ensure the losses are backpropagated correctly through:
a) the different models (autoencoder and inceptionv3)
b) the different branches of inceptionv3
Code
Here is my Lightning module code:
  def __init__(self, classCount, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.module = ae_cnn_models.AECNN(classCount)
    self.classCount = classCount
    # ... more initialization elided for simplicity
    self.loss1 = torch.nn.MSELoss(size_average=True)
    self.loss2 = torch.nn.BCELoss(size_average=True)

  def configure_optimizers(self):
    # ... elided for simplicity

  def train_dataloader(self):
    # ... elided for simplicity

  def training_step(self, batch, batch_idx):
    x, y = batch
    input = self.trans_train(x)
    input = input.to(dtype=torch.float32, device=self.device)
    varInput = torch.autograd.Variable(input)
    varTarget1 = torch.autograd.Variable(input)
    varTarget2 = torch.autograd.Variable(y)
    varOutput1, varOutput2 = self.module.forward(varInput)
    classifierOut1, classifierOut2 = varOutput2

    # autoencoder model loss:
    lossvalue1 = self.loss1(varOutput1, varTarget1)
    # weighting between main and aux branch of inception model:
    lossvalue2 = 0.8 * self.loss2(classifierOut1, varTarget2) + 0.2 * self.loss2(classifierOut2, varTarget2)  
    # weighting btween MSE and BCE respectively:
    loss = 0.1 * lossvalue1 + 0.9 * lossvalue2

    output = {
      'loss': loss,  # required
    }
    return output

def trans_train(self, x):
  # ...  some image transformations elided for simplicity

if __name__ == "__main__":
  classCount = 14
  model = LightningChestXrayCnnClassifierInceptionV3(classCount)
  trainer = pl.Trainer(max_epochs=15)
  trainer.fit(model)

When I look in the original source code, I see they simply do:
optimizer.zero_grad()
lossvalue.backward()
 optimizer.step()

which seems identical to what Pytorch Lightning would be doing. But again, what's not clear is how is this loss backpropagated correctly through:
a) the different models (autoencoder and inceptionv3)
b) the different branches of inceptionv3
?
Should I configure multiple / weighted losses differently to ensure the correct losses are backpropagated to the respective models/branches of models? Or is setting up a single loss value in the training_step sufficient?
What's your environment?

OS: MacOS
Packaging conda
Version 0.8.4