Selecting topk loss values using one pass through data and reduction='none'

Sina_Mokhtarzadeh: <p>Hi,</p>
<p>I’m trying to optimize using only a subset of the losses of the data in a batch. I basically run the model regularly with a big batch with no reduction on loss terms. Then based on some criteria I choose some of the individual loss therms. Then I take a mean of these and backward based on this loss. Is doing this ok?</p>
<p>a sample:</p>
<pre><code class="lang-auto">loss_all = run_model(batch_all)
loss_sub = torch.index_select(loss_all, 0, topk_idxs)
loss = loss_sub.mean()
loss.backward()
</code></pre>
<p>Is this ok? Or for such a use case I have to run with no_grad first, select data, then run model with the selected data?</p>
<p>Thanks</p>

Sina_Mokhtarzadeh: <p>Ok I did a toy example and this works like selecting the inputs from the beginning. However, the computational cost of backward doesn’t decrease. Is there a way to exclude unwanted samples of the batch from backward based on the loss after doing a forward with all the samples (to speed up the backward pass)?</p>

KFrank: <p>Hi Sina!</p>
<aside class="quote no-group" data-username="Sina_Mokhtarzadeh" data-post="2" data-topic="166071" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sina_mokhtarzadeh/48/54761_2.png" class="avatar"> Sina_Mokhtarzadeh:</div>
<blockquote>
<p>However, the computational cost of backward doesn’t decrease. Is there a way to exclude unwanted samples of the batch from backward</p>
</blockquote>
</aside>
<p>No, there isn’t.  The problem is that although in your use case, the batch<br>
elements are independent of one another – that is, their gradients don’t<br>
mix together – <code>backward()</code> doesn’t and can’t know this.</p>
<p>Suppose you had multiple channels that were mixed together in the forward<br>
pass and then you only wanted to compute the gradients of the <code>topk</code> largest<br>
<em>channel</em> losses.  <code>backward()</code> would have to process all of the channels to<br>
get the correct gradients – because they mix together.  But <code>backward()</code><br>
can’t tell the independent-batch-elements use case from the<br>
mixed-together-channels use case, so it (thinks it) has to process all of<br>
your batch elements even though you only want gradients for the batch<br>
elements with the <code>topk</code> largest losses.</p>
<aside class="quote no-group" data-username="Sina_Mokhtarzadeh" data-post="1" data-topic="166071">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sina_mokhtarzadeh/48/54761_2.png" class="avatar"> Sina_Mokhtarzadeh:</div>
<blockquote>
<p>Or for such a use case I have to run with no_grad first, select data, then run model with the selected data?</p>
</blockquote>
</aside>
<p>This approach has the potential to make your backward pass cheaper (but<br>
at the cost of performing the additional forward pass).  Also, you are now<br>
only performing the full-batch forward pass <em>without</em> tracking gradients, so<br>
this could also offer some cost savings.</p>
<p>However, depending on the structure of your model and how your gpu (or<br>
cpu) is being used, running your forward / backward pass on just your<br>
<code>topk</code> batch may or may not be significantly cheaper than running it on<br>
your full batch, so you’d have to time the two versions to be sure that such<br>
a scheme actually helps.</p>
<p>Best.</p>
<p>K. Frank</p>

Sina_Mokhtarzadeh: <p>Thanks KFrank for your comment.</p>
<p>Actually, I think the batch data is independent in my case (Not sure if batch norm breaks that independence though <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20">).<br>
I don’t want to do this on the other dimensions like channels that could mix.</p>
<p>Currently, I’m trying to fix this by modifying pack_hook and unpack_hook in saved_tensors. Inside unpack_hook I want to select only the desired batch indices. I’m not sure this is the way but it was the closest path that I found. I would be glad if anyone has a say on this.</p>
<p>Also, I don’t want to use torch.no_grad since the ratio between the selected subset and total batch size is high and this will end up taking more time than just running it without torch.no_grad.</p>