How to update parameters in pytorch with multiple loss functions?

user_17286074 0
Viewed 113

I have met some problems in updating parameters while training model. For instance, I have 2 loss-functions: loss1 and loss2, let loss = loss1 + loss2. The model has feature extractor layers and classifier layers. And I want to update the parameters in feature extractor layers using loss and the parameters in classifier layers using loss1. How can I do that in pytorch??? Thanks guys!

pytorch deep-learning
1 Answers

Let fe be the feature extractor layer and fc the classifier layer. One approach is to perform two backward passes:

  • the first one on loss2 with the fc layer gradients deactivated,

  • the second one on loss1 on both layers fc and fe.

As a result, fc should have the gradients relevant to loss1, while fe will have gradients from both loss1 and loss2, i.e. loss.

>>> fc.requires_grad_(False)
>>> loss2.backward(retain_graph=True)
>>> fc.requires_grad_(True)
>>> loss1.backward()

Do make sure this is the desired result you are looking for.

user_63313690
Related