Backwards with only a subset of the output losses?

username: Juuso_Korhonen
cooked: <p>Hi everyone,</p>
<p>I see that the feedforward of my neural network takes only half the amount of time as the backward (pretty common stat with neural networks).I would like to reduce the amount of backwards calculations by only doing them for outputs that achieve above average loss. How do I do this correctly? Here’s the code:</p>
<pre><code class="lang-auto">preds, truths, loss = self.train_step(data) # each are [4096,1] 
# Get only hard samples
masked_loss = loss[0:5].mean() # taking explicitly only 5 samples
self.scaler.scale(masked_loss).backward() 
self.scaler.step(self.optimizer) 
self.scaler.update()
</code></pre>
<p>But I noticed that the backward() operation is still taking the same time as just doing it for the whole loss. What am I missing?</p>

username: J_Johnson
cooked: <p>The increased time is likely due to the number of calculation steps involved. But calculations across the batch dim are performed asynchronuously.</p>
<p>Let me explain it this way. Suppose you’re using an H100 with 18,432 CUDA cores. And your batch size is 100. Each of those cores behave like an independent worker, capable of making quick work of tensor operations with multiplication of hundreds of elements in a matmul per batch sample.</p>
<p>Decreasing the batch size won’t necessarily decrease the speed to perform a given set of calculations needing to be performed in order, such as a backward pass.</p>
<p>But it would be interesting to see if training a model on only outliers has any benefit to accuracy. My initial inclination is it might result in overfitting.</p>

username: Juuso_Korhonen
cooked: <p><a class="mention" href="/u/j_johnson">@J_Johnson</a> Thanks for the answer. That makes sense. Actually this Online Hard Example Mining (OHEM) paper seems to suggest that focusing on the hard examples might lead to better test accuracy. They leave the reasons as to be investigated, and thats what I’m trying to do.</p>
<p>I actually found that in my particular problem the cost of feedforward is negligible compared to the backward calculations. So I devised this kind of plan that I run massive amount of samples in feedforward and then choose only a subset for both the feedforward and backward:</p>
<pre><code class="lang-auto">            start_time_forward = time.time()
            if self.opt.OHEM: # Online Hard Example Mining, data has 2*K samples
                # Massive forward pass (with 2*K samples), do not calculate gradients
                with torch.no_grad():
                    start_time_forward1 = time.time()
                    _,_,loss = self.train_step(data) # _,_, [2*K,1]
                    end_time_forward1 = time.time() # takes ~ 0.0025, negligible
                    # choose only K hard samples
                    start_time_masking = time.time()
                    data = self.mask_data(data, torch.topk(loss.detach(), K, sorted=False)[1])
                    end_time_masking = time.time() # takes ~ 5e-5, negligible
                # Mini forward pass (with K samples), calculate the gradients
                start_time_forward2 = time.time()
                preds, truths, loss = self.train_step(data) # data has now only K samples, so SHOULD TAKE THE SAME TIME AS REGULAR BELOW!!!
                end_time_forward2 = time.time() # takes ~ 0.06
                end_time_forward = time.time() # takes ~ 0.06, so the mini forward pass TAKES UP ALL THE TIME!!! TODO: solve why?
            else: # Regular forward pass with all samples, data has K samples
                preds, truths, loss = self.train_step(data)
                end_time_forward = time.time() # takes ~ 0.002, negligible
</code></pre>
<p>But as you can see I have something weird happening. That second feed forward in the OHEM case is taking 0.06 seconds, where as the regular feedforward is taking 0.002. They should be the same as they both run K samples. Am I doing some step wrong there?</p>

username: ptrblck
cooked: <p>CUDA operations are executed asynchronously so you would need to synchronize your code before starting or stopping your host timers.</p>

username: J_Johnson
cooked: <p>I suppose it might work if the outliers are in fact accurately labeled and you maintain a loss cutoff that allows a substantial amount of the samples through. But most large public datasets have mislabeled data. Something to watch out for.</p>
<p>This looks like a decent tutorial for code synchronization in Pytorch, as <a class="mention" href="/u/ptrblck">@ptrblck</a> raised:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://www.speechmatics.com/company/articles-and-news/timing-operations-in-pytorch">
  <header class="source">

      <a href="https://www.speechmatics.com/company/articles-and-news/timing-operations-in-pytorch" target="_blank" rel="noopener nofollow ugc">Speechmatics</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/361;"><img src="https://images.ctfassets.net/yze1aysi0225/47pOjt3KagFbO35ICiWW4F/2452c35aa4b836ba1385bdfed12e5bee/Timings_Pytorch_Blog_Blog_Header_Timings_Pytorch_Image_16-9.svg" class="thumbnail" width="" height=""></div>

<h3><a href="https://www.speechmatics.com/company/articles-and-news/timing-operations-in-pytorch" target="_blank" rel="noopener nofollow ugc">The future of AI tech learning &amp; development | Speechmatics</a></h3>

  <p>In a world of increasingly costly machine learning model deployments, ensuring accurate GPU operation timing is key to resource optimization. Read more!</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


username: Juuso_Korhonen
cooked: <p>Thanks <a class="mention" href="/u/j_johnson">@J_Johnson</a> and <a class="mention" href="/u/ptrblck">@ptrblck</a> for advising me using the synchronization. It revealed that the running times I calculated in my second answer were wrong and indeed the backwards takes 2x the time of feedforward also in my case (the standard rule of thumb).</p>
<p>Going back to the original question, I see that I can decrease the time it takes to do the backward if I do the forward pass two times (first with all samples and then with only hard samples):</p>
<pre><code class="lang-auto">            if self.opt.OHEM: # Online Hard Example Mining
                # 1st forward pass, do not calculate gradients
                with torch.no_grad():
                    _,_,loss = self.train_step(data) # _,_, [K,1]
                    torch.cuda.synchronize()
                    # choose only K hard samples
                    start_time_masking = time.time()
                    hard_fraction = 0.25 # reduction in computation time should be roughly: hard_fraction * 2/3
                    data = self.mask_data(data, torch.topk(loss.detach(), int(self.opt.num_rays * hard_fraction), sorted=False)[1])
                    end_time_masking = time.time() # takes ~ 5e-5, negligible
                # 2nd forward pass, calculate the gradients. Data now only has a &lt;&lt;hard_fraction&gt;&gt; of the samples
                preds, truths, loss = self.train_step(data)
                torch.cuda.synchronize()
                end_time_forward = time.time() # takes ~ 0.06
            else: # Regular forward pass with all samples, data has K samples
                preds, truths, loss = self.train_step(data)
                torch.cuda.synchronize()
                end_time_forward = time.time() # takes ~ 0.03

            loss = loss.mean()
            start_time_backward = time.time()
            # Backward pass
            self.scaler.scale(loss).backward()
            self.scaler.step(self.optimizer)
            self.scaler.update()
            torch.cuda.synchronize()
            end_time_backward = time.time() # takes ~ 0.02 with OHEM, and ~ 0.06 otherwise.
</code></pre>
<p>BUT I notice that if I do it this simpler way (with only one forward pass and selecting only a subset loss for backward):</p>
<pre><code class="lang-auto">            start_time_forward = time.time()
            preds, truths, loss = self.train_step(data) # data has K samples
            torch.cuda.synchronize()
            end_time_forward = time.time() # takes ~ 0.03 with K samples
            # OHEM: pick only the hard samples for the costy backward pass (backward pass cost = 2 x forward pass cost)
            if self.opt.OHEM:
                hard_fraction = 0.25 # reduction in computation time SHOULD be roughly: hard_fraction * 2/3
                loss = loss[torch.topk(loss.detach(), int(self.opt.num_rays * hard_fraction), sorted=False)[1]] # a subset loss

            loss = loss.mean() # turn loss to scalar
            start_time_backward = time.time()
            # Backward pass
            self.scaler.scale(loss).backward()
            self.scaler.step(self.optimizer)
            self.scaler.update()
            torch.cuda.synchronize()
            end_time_backward = time.time() # still takes ~ 0.06 even with the OHEM.
</code></pre>
<p>The time it takes to calculate the backward stays roughly the same as I wouldn’t do the OHEM picking of hard samples. This makes me think that its still doing the gradient calculations for all the samples. I guess it should be possible to gain performance, since we see that by doing it with two forward passes works on reducing the time it takes to calculate the backwards.</p>

username: KFrank
cooked: <p>Hi Juuso!</p>
<aside class="quote no-group" data-username="Juuso_Korhonen" data-post="6" data-topic="188313" 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/juuso_korhonen/48/63920_2.png" class="avatar"> Juuso_Korhonen:</div>
<blockquote>
<p>This makes me think that its still doing the gradient calculations for all the samples.</p>
</blockquote>
</aside>
<p>This is correct.  Pytorch operations, including those in the backward pass,<br>
operate, for efficiency reasons, on  entire tensors (even if they only use or<br>
modify a subset of elements).</p>
<p>Let me reword your code slightly:</p>
<pre data-code-wrap="python"><code class="lang-python">loss_subset = loss[torch.topk(loss.detach(), int(self.opt.num_rays * hard_fraction), sorted=False)[1]] # a subset loss

loss_scalar = loss_subset.mean() # turn loss to scalar
</code></pre>
<p>(This is just to make clear that in your code, <code>loss</code> is a python name that refers<br>
to different tensors at different points in the code.)</p>
<p>When you call <code>loss_scalar.backward()</code>, you backpropagate through the scalar<br>
back to <code>loss_subset</code> – a tensor that does consist of only a subset of batch items.<br>
This very trivial backpropagation is ever so slightly cheaper because <code>loss_subset</code><br>
is a smaller tensor (but any benefit is so small as to be irrelevant).</p>
<p>Autograd then backpropagates from <code>loss_subset</code> back to <code>loss</code>.  (This is also a<br>
trivial backpropagation.)  However, <code>loss</code> now contains <em>all</em> of the batch elements,<br>
even though the gradients being backpropagated for many of those batch elements<br>
are zero.</p>
<p>All of the rest of the backpropagation is then carried out with <em>all</em> of the batch<br>
elements (because pytorch is processing entire tensors).  Even where the gradient<br>
being backpropagated has zeros for the masked-out rows, pytorch still performs<br>
the entire tensor operation, multiplying (or whatever) by those zeros.  So there is<br>
no savings in time (nor in memory).</p>
<aside class="quote no-group">
<blockquote>
<p>I guess it should be possible to gain performance,</p>
</blockquote>
</aside>
<p>Because of how pytorch operates on entire tensors, there is no way to get your<br>
hoped-for performance gain if you only perform the single forward pass.</p>
<aside class="quote no-group">
<blockquote>
<p>since we see that by doing it with two forward passes works on reducing the time it takes to calculate the backwards.</p>
</blockquote>
</aside>
<p>Your scheme of performing two forward passes it the way to achieve your<br>
performance gain.  You <em>need</em> to perform the first forward pass in order to<br>
compute the per-batch-element losses that you use to perform your “hard<br>
mining.”  But if you want to backpropagate for only the “hard” subset of your<br>
batch elements, you also <em>need</em> to perform the second forward pass to<br>
construct the computation graph that contains the smaller hard-subset<br>
intermediate tensors.</p>
<p>If <code>hard_fraction</code> is close to zero, you will likely get significant savings, while<br>
if <code>hard_fraction</code> is close to one, the cost of the second forward pass will<br>
likely exceed any savings from the subset backward pass.</p>
<p>As an aside, wrapping your first forward pass in a <code>with torch.no_grad():</code><br>
block will get you some additional savings, mostly in memory, because doing<br>
so avoids constructing the full, non-subset computation graph that you won’t<br>
be using anyway.  (These memory savings are likely to be significant.)</p>
<p>Last, let me second J’s comment that your hard mining may or may not improve<br>
training speed and / or final performance.  Trying it both ways, of course, is the<br>
best way to find out.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Juuso_Korhonen
cooked: <p>Thank you <a class="mention" href="/u/kfrank">@KFrank</a>. Alright the (only) way to do it with PyTorch is with 2 forward passes. Is this specific to PyTorch and to the way they do the backpropagation with the graph? Do you know if others, like tensorflow or JAX, would differ in this regard?</p>