How to Detach specific components in the loss?

username: ankitvad
cooked: <p>I’m a little confused how to detach certain model from the loss computation graph.<br>
If I have 3 models that generate an output: <code>A</code>,  <code>B</code> and <code>C</code>, given an input.<br>
And, 3  optimizers for those 3 models: <code>Oa</code>, <code>Ob</code> and <code>Oc</code>, for A, B and C resp.<br>
Let’s assume I have an initial input <code>x</code>.<br>
I also have 3 losses <code>L1</code>, <code>L2</code> and <code>L3</code>.</p>
<p>Let’s assume a scenario where the system does <code>L2</code> that calculates <code>B(A(x))</code>, amongst other things.</p>
<p>If we want a loss to optimize only certain models. eg. If we want to optimize only <code>B</code> we can do it 2 ways:</p>
<pre><code class="lang-auto">1)
op = A(x)
op2 = B(op.detach())
cost = L2(op2, ground_truth)
cost.backward()
Ob.step()
#This would only backprop. gradients and optimize B.

2)
op = A(x)
op2 = B(op)
cost = L2(op2, ground_truth)
cost.backward()
Ob.step()
#This would backprop. loss to both A and B, 
#but it doesn't matter because we are just doing Ob.step()
</code></pre>
<p>Now in this if we wanted to optimize <code>A</code> for <code>B(A(x))</code> we can’t do <code>(1)</code> we can do <code>(2)</code>.</p>
<pre><code class="lang-auto">op = A(x)
op2 = B(op)
cost = L2(op2, ground_truth)
cost.backward()
Oa.step()
#Here gradients are calculated for both `A` and `B`, but it doesn't matter because we are
#just optimizing A with Oa.
</code></pre>
<p>I have understood this but let’s say that:<br>
netLoss = L1 + L2 + L3<br>
In this I have the second case above. L2 has both A and B in the computation graph and we only want to optimize A with respect to L2. The problem is that we want to do netLoss.backward() and L1 + L2 + L3 contains A, B and C in them total.</p>
<p>So after netLoss.backward() we have to do:<br>
Oa.step()<br>
Ob.step()<br>
Oc.step()</p>
<p>So, <code>Ob</code> will take gradients from L2 as well, and form the other losses that has it also. Is there some way to do netLoss.backward() when we can exclude certain models from the graph? Eg. in (2) above where we want to update only <code>A</code> in B(A(x)).</p>

username: ankitvad
cooked: <p><a class="mention" href="/u/apaszke">@apaszke</a> <a class="mention" href="/u/simonw">@SimonW</a> <a class="mention" href="/u/alband">@albanD</a></p>

username: apaszke
cooked: <p>The information we would need to give you the best possible schedule includes how those losses are exactly influenced by individual modules.</p>

username: ankitvad
cooked: <p><a class="mention" href="/u/apaszke">@apaszke</a>  Please consider the Eq. 8, 9, 10, 11. in <a href="https://arxiv.org/pdf/1711.06861.pdf" rel="nofollow noopener">https://arxiv.org/pdf/1711.06861.pdf</a><br>
Eq. 8 - Minimize CrossEntropy Loss<br>
Eq 9 - Maximize Entropy of Label(minimizing the negative of Entropy)<br>
which is a different Loss not present in the package but can be created as: <code>-∑ p.log(p)</code> and i’ll have to manually construct it.<br>
Eq 10 - Minimize CrossEntropy Loss</p>
<p>The issue is with Eq 9, where the computation is like: op = Classifier(Encoder(x)) and the Loss only wants to optimize the Encoder model for the <code>op</code> loss.</p>
<p>Like i mentioned in the example above, this would be easy to do if the loss was backpropogated individually, I could just do Encoder_Optimizer.step()</p>
<p>But, the paper model is end-to-end trained which means that they are just backpropogating Eq. 11; where Eq 11 = Eq 8 + Eq 9 + Eq 10</p>
<p>The problem here is that since some other loss ie. Eq 8 wants to optimize the Classifier model so the Classifier will end up being updated by gradients from both Eq 8 and Eq 9.</p>
<p>I hope i’m understanding it correctly, and, Is there some way to do this?</p>

username: apaszke
cooked: <p>Doing multiple backward passes before calling any optimizers will give you end to end training too. It’s just a name for joint optimization. I didn’t have time to look into the paper, but it’s possible that you could backprop from Eq9 loss first, zero grads of Classifier, and then backprop from Eq8 + Eq10 (remember that .grad is accumulated instead of overwritten).</p>

username: ankitvad
cooked: <p>Ahhh… I did not know that we could backprop. losses individually as well. Yes, this makes perfect sense. Thanks a lot. This really helped me out.</p>

username: ankitvad
cooked: <p><a class="mention" href="/u/apaszke">@apaszke</a> <a class="mention" href="/u/ptrblck">@ptrblck</a>  I just had a small follow up, little different query.<br>
If model1 gives an output a which is then given to model2 and maybe model3 also later on and we want to make sure that different models are updated based on different losses then how does .detach() work.</p>
<p>Let’s say a.detach() is the input to model2 so if we calculate loss and do .backward() gradients will only be calculated for model2 not model1.</p>
<p>But then in the next step when a is given to model3 and the new loss is done .backward() on does the gradient get calculated for model3 as well as model1 or just model3?</p>
<p>Basically, does calling .detach() return a copy of the original output ‘a’ or does it detach the original a from computation chain?</p>

username: ptrblck
cooked: <p><code>.detach()</code> will return a detached version of your tensor. You can reuse the “attached” tensor for further computations as long as you don’t reassign it somehow:</p>
<pre><code class="lang-python">modelA = nn.Linear(10, 10)
modelB = nn.Linear(10, 10)
modelC = nn.Linear(10, 10)

x = torch.randn(1, 10)
a = modelA(x)
b = modelB(a.detach())
b.mean().backward()
print(modelA.weight.grad)
print(modelB.weight.grad)
print(modelC.weight.grad)

c = modelC(a)
c.mean().backward()
print(modelA.weight.grad)
print(modelB.weight.grad)
print(modelC.weight.grad)
</code></pre>

username: ankitvad
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thanks a lot. This was really helpful <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=6" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: tjoseph
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> What about if modelC = modelB and I want to backpropagate through modelB, but optimize modelA?</p>
<pre><code class="lang-auto">modelA = nn.Linear(10, 10)
modelB = nn.Linear(10, 10)

# a is needed for both losses
x = torch.randn(1, 10)
a = modelA(x)

b1 = modelB(a.detach())
loss1 = loss_function_1(b1)
loss1.mean().backward()
# Here call optimForB.step()

b2 = modelB(a)
loss2 = loss_function_2(b2)
loss2.mean().backward()
# Here call optimForA.step()
</code></pre>
<p>I will need to compute</p>
<pre><code class="lang-auto">b1 = modelB(a.detach())
</code></pre>
<p>for the first loss and</p>
<pre><code class="lang-auto">b2 = modelB(a)
</code></pre>
<p>for the second. This seems kinda wasteful imo. Is there a way to detach after calculation, so I only have to calculate <code>b = modelB(a)</code> once?</p>

username: ptrblck
cooked: <p>I don’t see another approch, if the loss functions are different (in you current code snippet they are the same, so you could just use a single forward pass, but I assume that’s a typo).<br>
For the second forward pass, you could set the <code>requires_grad</code> attributes of all parameters in <code>modelB</code> to <code>False</code>, if you don’t need these gradients.</p>

username: tjoseph
cooked: <p>Okay, thanks. Good to know!</p>