How can i process multi loss in pytorch?
First and 3rd attempt are exactly the same and correct, while 2nd approach is completely wrong.
In Pytorch, low layer gradients are Not "overwritten" by subsequent backward()
calls, rather they are accumulated, or summed. This makes first and 3rd approach identical, though 1st approach might be preferable if you have low-memory GPU/RAM (a batch size of 1024 with one backward() + step()
call is same as having 8 batches of size 128 and 8 backward()
calls, with one step()
call in the end).
To illustrate the idea, here is a simple example. We want to get our tensor x
close to 40,50 and 60
simultaneously:
x = torch.tensor([1.0],requires_grad=True)
loss1 = criterion(40,x)
loss2 = criterion(50,x)
loss3 = criterion(60,x)
Now the first approach: (we use tensor.grad
to get current gradient for our tensor x
)
loss1.backward()
loss2.backward()
loss3.backward()
print(x.grad)
This outputs : tensor([-294.])
(EDIT: put retain_graph=True
in first two backward
calls for more complicated computational graphs)
The third approach:
loss = loss1+loss2+loss3
loss.backward()
print(x.grad)
Again the output is : tensor([-294.])
2nd approach is different because we don't call opt.zero_grad
after calling step()
method. This means in all 3 step
calls gradients of first backward
call is used. For example, if 3 losses provide gradients 5,1,4
for same weight, instead of having 10 (=5+1+4), now your weight will have 5*3+1*2+4*1=21
as gradient.
For further reading : Link 1,Link 2
Third attempt is the best.
Two different loss functions
If you have two different loss functions, finish the forwards for both of them separately, and then finally you can do (loss1 + loss2).backward()
. It’s a bit more efficient, skips quite some computation.
Extra tip: Sum the loss
In your code you want to do:
loss_sum += loss.item()
to make sure you do not keep track of the history of all your losses.
item()
will break the graph and thus allow it to be freed from one iteration of the loop to the next. Also you could use detach()
for the same.
-- Comment on first approach removed, see other answer --
Your second approach would require that you backpropagate with retain_graph=True
, which incurs heavy computational costs. Moreover, it is wrong, since you would have updated the network weights with the first optimizer step, and then your next backward()
call would compute the gradients prior to the update, which means that the second step()
call would insert noise into your updates. If on the other hand you performed another forward()
call to backpropagate through the updated weights, you would end up having an asynchronous optimization, since the first layers would be updated once with the first step()
, and then once more for each subsequent step()
call (not wrong per se, but inefficient and probably not what you wanted in the first place).
Long story short, the way to go is the last approach. Reduce each loss into a scalar, sum the losses and backpropagate the resulting loss. Side note; make sure your reduction scheme makes sense (e.g. if you are using reduction='sum' and the losses correspond to a multi-label classification, remember that the number of classes per objective is different, so the relative weight contributed by each loss would also be different)