Hi, this is my first post here and I’m not very familiar with pytorch yet. My situation is as follows:
I’m training a simple (sequential, dense) neural network where the loss consists of a sum of two terms, say:
loss = loss1 + loss2. I would like to investigate the contributions to the gradient of the network weights of each loss term separately (during the training). Unfortunately this does not seem to be straightforward, i.e. something along the lines of the following does not work:
loss1.backwards()
grad1 = model.layername.weight.grad.detach().clone()
loss2.backwards()
grad2 = model.layername.weight.grad.detach().clone()
I’ve also tried setting retain_graph=True as well as optimizer.zero_grad() in between.
However I either get that grad1 is exactly equal to grad2, which makes no sense since loss1 and loss2 are obtained through completely different functions, or I get that grad2 is just a zero-tensor (which also does not make any sense since the optimization w.r.t. the total loss loss yields very different results that the optimization w.r.t. loss1 alone).
What am I doing wrong, resp. how can I obtain the individual gradients w.r.t. to each summand of the total loss function in order to analyze/compare then?
Thank you very much for any help!
Your code snippet should generally work (besides the typo in backward() as seen here:
# setup
model = nn.Linear(1, 1, bias=False)
x = torch.randn(1, 1)
# zero out grads
out = model(x)
loss1 = out.mean()
loss2 = ((out - 1)**2).mean()
loss1.backward(retain_graph=True)
grad1 = model.weight.grad.clone()
model.zero_grad()
loss2.backward()
grad2 = model.weight.grad.clone()
print(grad1)
# tensor([[-1.9150]])
print(grad2)
# tensor([[4.5179]])
# clone each grad
model.zero_grad()
out = model(x)
loss1 = out.mean()
loss2 = ((out - 1)**2).mean()
loss1.backward(retain_graph=True)
grad1 = model.weight.grad.clone()
loss2.backward()
grad2 = model.weight.grad.clone()
print(grad1)
# tensor([[-1.9150]])
print(grad2)
# tensor([[2.6030]]) # = 4.5179 + (-1.9150)
# single loss
model.zero_grad()
out = model(x)
loss1 = out.mean()
loss2 = ((out - 1)**2).mean()
loss = loss1 + loss2
loss.backward()
print(model.weight.grad)
# tensor([[2.6030]])
Thank you very much for your answer and the example code @ptrblck!
This helped me greatly in figuring out what the issue was in my code. In fact, I accidentally had a .detach() in part of the computation of loss2, that’s why it’s gradient vanished (however it still had an effect on other parts of the network, which is presumably why I got very different results when training w.r.t. loss and loss1).
In case I’ll ever encounter a similar situation: Is there a way to analyse the effect of the back-propagation step by step in order to ‘debug’ or detect undesired behavior (such as disconnectedness of the graph)?
It can be quite difficult to detect such bugs if the loss function is a complicated custom function (as in my case).
I very much appreciate you taking the time to answer my question!
Good to hear it’s working now!
To debug the computation graph you could visualize it via e.g. torchviz.
Great, I’ll take a look at torchviz. Thanks again!