Is a sub-batch-tensor based on a discrete selection operation still differentiable when passed through multiple models?

ankitvad: <p>If I have 3 sub-models in my end-to-end pipeline - A, B and C! And, I have an input of Batch Size of - 5 (any number - 5 for convenience) = IP.</p>
<p>Let’s say - I have 2 steps of operations:</p>
<ol>
<li>
<em>(IP)[5 x N]</em> → <strong>A</strong> → <em>OP1</em>[5 x H] → <strong>B</strong> → OP2 [5 x 2]<br>
Here the input is sent through Model A and the output is then sent to Model B which gives a [BatchSize x 2] Output which can be used to obtain a Boolean Tensor of size BatchSize; let’s say : <em>OP2 = [0,1,0,1,0]</em>
</li>
</ol>
<p>Now, I want to use this tensor - <em>OP2</em> to select a sub-tensor from <em>OP1</em>, which comes out of model A.</p>
<pre><code class="lang-auto">OP1_SubTensor = OP1[OP2.nonzero(), :]
</code></pre>
<ol start="2">
<li>Ideally, I wanna use the OP1_SubTensor as input for Model C.</li>
</ol>
<p><em>OP1_SubTensor</em> [2 x H] → <strong>C</strong> → <em>OP3</em> [2 x M]</p>
<p>I have 2 losses! One on the output of model B - Loss1 and one on the output of model C - Loss 2.</p>
<p>I am wondering if I can use this discrete selection on the output of model A (OP1) and then pass that sub-tensor to model C and then when I back prop on Loss1 and Loss 2; It will all be end-to-end optimized!</p>
<p>Or, should I filter the <em>IP</em> [5 x N] based on <em>OP2</em> boolean tensor and pass it through the model A again?</p>
<p>Kinda confused!</p>

ankitvad: <p>So - upon checking other questions here and building on 2 replies -</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> <a href="https://discuss.pytorch.org/t/how-to-detach-specific-components-in-the-loss/13983/8" class="inline-onebox">How to Detach specific components in the loss? - #8 by ptrblck</a></p>
<p>and,<br>
<a class="mention" href="/u/alband">@albanD</a> <a href="https://discuss.pytorch.org/t/is-loss-backward-function-calculate-gradients-over-mini-batch/66873/5" class="inline-onebox">Is Loss.backward() function calculate gradients over mini-batch - #5 by albanD</a></p>
<p>It seems there are 2 possible approaches?</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

BS = 5 #batch size
modelA = nn.Linear(5,10)
modelB = nn.Linear(10, 2)
modelC = nn.Linear(10, 20)
x = torch.randn(BS, 5)

#Approach 1 - Clone to duplicate the computation Graph!
a = modelA(x)
a_copy = a.clone()
b = modelB(a_copy)
filter_index = torch.topk(b,1)[1].squeeze() #tensor([0, 1, 1, 1, 0])
a = a[filter_index.nonzero(),:].squeeze(1)
c = modelC(a)

print(modelA.weight.grad)
print(modelB.weight.grad)
print(modelC.weight.grad)

b.mean().backward(retain_graph = True)
#print(modelA.weight.grad)
#print(modelB.weight.grad)
#print(modelC.weight.grad)

c.mean().backward()
#print(modelA.weight.grad)
#print(modelB.weight.grad)
#print(modelC.weight.grad)


#Clear All:
modelA.weight.grad = None
modelB.weight.grad = None
modelC.weight.grad = None


#Approach 2 - Just retain_graph so it is not removed?
a = modelA(x)
b = modelB(a)
filter_index = torch.topk(b,1)[1].squeeze() #tensor([0, 1, 1, 1, 0])
a = a[filter_index.nonzero(),:].squeeze(1)
c = modelC(a)

print(modelA.weight.grad)
print(modelB.weight.grad)
print(modelC.weight.grad)

b.mean().backward(retain_graph = True)
#print(modelA.weight.grad)
#print(modelB.weight.grad)
#print(modelC.weight.grad)

c.mean().backward()
#print(modelA.weight.grad)
#print(modelB.weight.grad)
#print(modelC.weight.grad)
</code></pre>
<p>And, upon gradient checking, both seem to work fine! After the first <code>.backward()</code>, <em>modelA</em> and <em>modelB</em> is updated and the second <code>.backward()</code> updates <em>modelA</em> (further!) and <em>modelC</em>?</p>
<p>So, is there a potential difference between these 2 approaches? It seems - both require me to specify <code>retain_graph</code>. Unless I do something like this -</p>
<pre><code class="lang-auto">#Approach 3:
loss = b.mean() + c.mean()
loss.backward()
</code></pre>
<p>And this seems to work as well?! Without the <code>retain_graph</code> since the computation graph is combined (maybe?).</p>
<p>So - basically - Are all 3 approaches the same? and neither is better than the other? And, the <code>a = a[filter_index.nonzero(),:].squeeze(1)</code> step to discretely and partially filter the <em>modelA</em> output-batch poses no issue in breaking the computation graph?</p>
<p>PS: I checked this reply about memory efficiency using batch gradients accumulation here - <a href="https://discuss.pytorch.org/t/using-sub-batches-to-avoid-busting-the-memory/66772/2" class="inline-onebox">Using sub-batches to avoid busting the memory? - #2 by albanD</a> and if multiple batches can be evaluated (<code>.backward()</code>) for a single optimizer-step to accumulate gradients then a part of the batch output can be used to evaluate the model as well right?</p>