Autograd for different loss functions in different layers

Tiago: <p>Hey all,</p>
<p>I am having a problem with the autograd for different losses in diffrent layers, I want to implement a new algorithm, but first I did a check test to see if results are the same.</p>
<pre><code class="lang-auto">import torch 
import torch.nn as nn

device = 'cpu'

#define neural network
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Linear(1, 1, bias=False))
        self.layer2 = nn.Sequential(
            nn.Linear(1, 1, bias=False))
        
    def forward(self, x):
        out1 = self.layer1(x)
        out2 = self.layer2(out1)
        return out1,out2
    
model = Net().to(device)

#desired output
t = (torch.randn(1, 1)).to(device)

#input
inp = (torch.randn(1, 1)).to(device)

#run network
outputs = model(inp)

#loss from the ouput layer
loss = 0.5*torch.sum( (t - outputs[1])**2.0)

#calculate the gradients just to check
loss.backward(retain_graph=True)

#gradient of the second weight
grad1 = torch.autograd.grad(loss, model.layer2[0].weight, retain_graph=True)

#part of the chain rule to calculate the gradient of the first weight
t1 = torch.autograd.grad(-loss, outputs[0], create_graph=True, retain_graph=True)[0]

#so it will be same as backpropagation
t1 += outputs[0]

#new loss for the gradient
loss2 = 0.5*torch.sum( (t1 - outputs[0])**2.0)

#gradient for the first weight
grad2 = torch.autograd.grad(loss2, model.layer1[0].weight)

print(grad2[0])
print(model.layer1[0].weight.grad)
</code></pre>
<p>loss2 is the loss function for the hidden layer. Mathematically this should be the same as the usual backpropagation, the results should be equal, but they aren’t. I checked t1 and it is all right, but something is going on with grad2, can someone help me with this?</p>

albanD: <p>Hi,</p>
<p>Could you explain why grad2 should be the same as plain gradient? it uses gradient penalty computed in a differentiable manner. meaning that you introduce second order gradients here. Maybe you did not want the <code>create_graph=True</code>?</p>

Tiago: <p>Sure,</p>
<p>In the usual way, the gradient of the first weight would be:<br>
grad1 = -(t - outputs[1]) * outputs[1]<br>
grad2 = -(t-outputs[1])* model.layer2[0].weight * outputs[0]</p>
<p>For the new loss function in the hidden layer, I have:<br>
L = 0.5*torch.sum( (t1 - outputs[0])*<em>2.0)<br>
and t1 is:<br>
t1 = outputs[0] + (t-outputs[1])</em> model.layer2[0].weight<br>
Which the algorithm is right</p>
<p>Now, I want the gradient for the first weight is the derivative of L with respect to the weight, which is:<br>
grad = -(t1 - outputs[0]) * outputs[0]<br>
grad = -(t-outputs[1]) * model.layer2[0].weight * outputs[0]</p>
<p>Which is same as the usual way, if I disable create_graph=True, then I get zero gradient, maybe I am understanding wrong this create_graph?</p>

albanD: <p>By <code>grad1</code> you mean <code>model.layer2[0].weight.grad</code> wrt to the first loss you computed ?<br>
If so,the loss is actually <code>loss = 0.5*torch.sum( (t - outputs[0] * model.layer2[0].weight)**2.0)</code><br>
So <code>grad1</code> would be <code>-(t - outputs[1]) * outputs[0]</code> and not <code>-(t - outputs[1]) * outputs[1]</code>.</p>
<p>And if <code>grad2</code> is for <code>model.layer1[0].weight.grad</code> wrt to the first loss computed, this loss is <code>loss = 0.5*torch.sum( (t - x * model.layer1[0].weight * model.layer2[0].weight)**2.0)</code> and so you get <code>grad2 = -(t-outputs[1])* model.layer2[0].weight * x</code> right?</p>

Tiago: <p>Yes, grad1 is correct, I made a typo.</p>
<p>For grad2, that indeed means <code>model.layer1[0].weight.grad</code> I thought I could write directly the output values, without the weights, in the loss  function, since I would like to use t1 instead of t, but this is not possible, is it?</p>

albanD: <p>I don’t think it is.<br>
Also why do you compute <code>t1</code> with <code>-loss</code>?</p>

Tiago: <p>Ah too bad, I will have to find I workaround then, because I really need the derivative of the loss with t1 present there.</p>
<p>The minus sign is because I wrote the loss as (t - output) instead of (output - t), the update of the weights will be w += lr * grad instead of w += -lr * grad later on.</p>
<p>But thank you anyway for the help.</p>

albanD: <blockquote>
<p>The minus sign is because I wrote the loss as (t - output) instead of (output - t)</p>
</blockquote>
<p>But this is squared. So both are exactly the same. You don’t need a <code>-</code> sign here right?</p>

Tiago: <p>Sorry, I meant the derivative of the loss.</p>