Totally different training behaviour for 2 similar scenarios

username: yuri
cooked: <pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
dtype = torch.float32

X = torch.tensor([[1, 2, 3, 4, 5, 6]], dtype=dtype)
Y = torch.tensor([[1, 4, 9, 16, 25, 36]], dtype=dtype) 

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        torch.manual_seed(3)
        self.base_l1 = torch.nn.Linear(6, 6, bias=True)
        self.base_l2 = torch.nn.Linear(6, 6, bias=True)
        self.l3 = torch.nn.Linear(6, 6, bias=True)
        self.l4 = torch.nn.Linear(6, 6, bias=True)

    def forward(self, x):
        x1 = self.base_l1(x)
        x1 = F.relu(x1)
        x1 = self.base_l2(x1)
        x2 = x1#.detach()
        x2 = F.relu(x2)
        x2 = self.l3(x2)
        x2 = F.relu(x2)
        x2 = self.l4(x2)
        return x2, x1


model = Model()
my_list = ['base']
base_params = list(filter(lambda kv: my_list[0] in kv[0], model.named_parameters()))
params = list(filter(lambda kv: my_list[0] not in kv[0], model.named_parameters()))
prms = []
base_prms = []
for i in params:
    prms.append(i[1])
for i in base_params:
    base_prms.append(i[1])

# these 2 optimizers does not have any common parametrs
optimizer1 = optim.SGD(base_prms, lr=.05, momentum=0.9)
optimizer2 = optim.SGD(prms, lr=.05, momentum=0.9) 

Loss = nn.MSELoss()
l1 = []
l2 = []
n_iters = 300
for epoch in range(n_iters):
    print(epoch)
    optimizer1.zero_grad()
    optimizer2.zero_grad()
    y_pred2, y_pred1 = model(X)
    loss2 = Loss(y_pred2, Y)

    l2.append(loss2)
    loss2.backward()
    optimizer2.step()
    # a trick to avoid the RuntimeError: Trying to backward through the graph a second 
    # time ...
    y_pred1 = torch.tensor(y_pred1, requires_grad=True)
    loss1 = Loss(y_pred1, Y)
    l1.append(loss1)
    loss1.backward()
    optimizer1.step()

</code></pre>
<p>**<br>
by running this version we observe a very poor performance and we could even argue that y_pred1 part is not getting any closer to the target Y. However, if we rerun the same version by assigning to x2= x1.detach() in the forward function and cancelling the mentioned trick of restarting the y_pred1 tensor we get a very satisfying accuracy.<br>
I could also join the gradient graph of this model you could easily notice that l1_base and l2_base layers can be learned without any help from l3 and l4 and vice versa if we backpropagate separately from the green rectangles presenting y_pred2 (the lower one) and y_pred1 (the upper one)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b4021268efe63b5420a6f881b41f01e0c3ad420.png" data-download-href="https://discuss.pytorch.org/uploads/default/7b4021268efe63b5420a6f881b41f01e0c3ad420" title="Screenshot (23)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7b4021268efe63b5420a6f881b41f01e0c3ad420_2_384x500.png" alt="Screenshot (23)" data-base62-sha1="hAk7N5zJnj6wzhlQOLBLxEVjxok" width="384" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/b/7b4021268efe63b5420a6f881b41f01e0c3ad420_2_384x500.png, https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b4021268efe63b5420a6f881b41f01e0c3ad420.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b4021268efe63b5420a6f881b41f01e0c3ad420.png 2x" data-dominant-color="ECEFEE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot (23)</span><span class="informations">519×675 34.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: ptrblck
cooked: <p>This “trick” won’t work as you are breaking the computation graph:</p>
<pre><code class="lang-python">    # a trick to avoid the RuntimeError: Trying to backward through the graph a second 
    # time ...
    y_pred1 = torch.tensor(y_pred1, requires_grad=True)
    loss1 = Loss(y_pred1, Y)
    l1.append(loss1)
    loss1.backward()
    optimizer1.step()
</code></pre>
<p>The <code>RuntimeError</code> is real and you could:</p>
<ul>
<li>accumulate the losses and call <code>backward()</code> once</li>
<li>call <code>backward()</code> on each loss separately and then <code>optimizer.step()</code>
</li>
<li>use <code>backward(retain_graph=True)</code> in the first <code>backward</code> call and make sure the forward activations are not stale to compute the gradients for <code>loss1</code>. I would consider this use case an advanced method and would not recommend to use it unless you know exactly that you want this behavior.</li>
</ul>
<p>In your current code snippet you re avoiding the <code>RuntimeError</code> by detaching <code>y_pred1</code> from the computation graph by recreating a new tensor with <code>requires_grad=True</code>. While this won’t raise any errors, no gradients will be created in the model parameters since <code>y_pred1</code> is a new leaf tensor without any gradient history.</p>

username: yuri
cooked: <p>thank you for your reply and explanation. In fact, I had to go through all this because I need to update the value of y_pred1 based on the parameters update of optimezer2.step before running loss1. All this without altering the computation graph of y_pred1.  more concretely optimizer2.step()  will update the l3 layer. a column of this layer I want to reinject to y_pred1 data and then calculate loss1 based on this new value but with the same gradient history.<br>
<strong>thus I cannot use</strong> the <strong>first</strong> and <strong>second</strong> alternatives that you proposed.<br>
I am doing this:</p>
<pre><code class="lang-auto">       loss2.backward(retain_graph=True)
       optimizer2.step()
       y_pred1.data = model.l3.weight[0, :].data.unsqueeze(0) # I am trying this one
       l1.append(loss1)
       loss1.backward()
       optimizer1.step()
</code></pre>
<p>is this valid? am I assigning a new value to the tensor y_pred1 and preserving everything else about its gradient?<br>
why should y_pred1 be freed when calling loss2.backward() when the gradient is not a function of neither y_pred1, base_l1 nor base_l2? if retaining the gradient graph is a pain that we must go through could we do this partially? only for a set of parameters and not for the whole model?<br>
thank you</p>

username: ptrblck
cooked: <p>Using the <code>.data</code> attribute is not recommended, as Autograd won’t track this change.<br>
Based on your description, it seems however as if that’s exactly your use case. I.e. you want to inject a custom value into the graph, which was not created in a differentiable way.<br>
Given that, I would highly recommend to manually check the gradient calculation and make sure that the gradients are indeed calculated as expected.</p>
<aside class="quote no-group" data-username="yuri" data-post="3" data-topic="137775">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/y/71c47a/48.png" class="avatar"> yuri:</div>
<blockquote>
<p>why should y_pred1 be freed when calling loss2.backward() when the gradient is not a function of neither y_pred1, base_l1 nor base_l2?</p>
</blockquote>
</aside>
<p>I don’t know where this approach is coming from so could you explain it a bit more?</p>

username: yuri
cooked: <p>exactly! that is what I was trying to do! and as you mentioned a manual check of this operation is a must.<br>
about the second part of the question I think now I can formulate it in a better way:<br>
to avoid grpah_retain=True, especially for big size models, it is possible to run loss2.backward() only on a subset of the model’s leaves. I think this helps us avoid allocating memory for retain_grad &amp; avoid the wasted effort of calculating their gradient and then setting them back to 0 with optimizer2.zero_grad().</p>

username: ptrblck
cooked: <p>Yes, you can use the <code>inputs</code> argument to <code>loss.backward</code> to specify which tensors should be used for the gradient computation. All other tensors will be ignored.</p>