Prevent loss divergence during training while using AMP

Tariq: <p>Hello all,</p>
<p>I recently added AMP support for my code training a segmentation model with a resnet-50 backbone.<br>
Everything seems to be working well at the beginning, except that the loss diverges after some iterations.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/a/cabee50fe899a5327536d732073191c8289fc38b.png" data-download-href="https://discuss.pytorch.org/uploads/default/cabee50fe899a5327536d732073191c8289fc38b" title="losses"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/cabee50fe899a5327536d732073191c8289fc38b_2_690x218.png" alt="losses" data-base62-sha1="sVzreh7yLyBeMNe0KLdPz5yqm15" width="690" height="218" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/cabee50fe899a5327536d732073191c8289fc38b_2_690x218.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/cabee50fe899a5327536d732073191c8289fc38b_2_1035x327.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/a/cabee50fe899a5327536d732073191c8289fc38b_2_1380x436.png 2x" data-dominant-color="FAFAFA"><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">losses</span><span class="informations">1930×610 56.7 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>
<p>The orange curve corresponds to the training loss w.o mixed precision, the grey one uses mixed precision with the default parameters and the purple one has <code>init_scale</code> set to 8192.</p>
<p>I am not sure which parameters I should tweak in order to get it to work correctly and wasn’t able to find an explanation of the effect of the <code>GradScaler</code> parameters online.</p>
<p>What could cause this divergence and how can I prevent it ?</p>
<p>Thanks,</p>

ptrblck: <p>Mixed-precision training should not cause any divergence and should be able to reach a similar FP32 accuracy metric.<br>
Could you describe your use case a bit more and what the curves represent?<br>
It’s especially interesting since ResNet50 is being trained in mixed-precision for years without seeing failures.<br>
Also, which GPU and <code>dtype</code> (<code>float16</code> or <code>bfloat16</code>) are you using?</p>

Tariq: <p>Hello,</p>
<p>I am training a light-weight refinenet with a resnet-50 backbone for simultaneous segmentation (20 channels/classes) and regression (one additional channel).</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/DrSleep/light-weight-refinenet">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/DrSleep/light-weight-refinenet" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/344;"><img src="https://opengraph.githubassets.com/b8f83ae907e4b615da2ceeb8e119992c24a91b3c064418e8634c8ce966224986/DrSleep/light-weight-refinenet" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/DrSleep/light-weight-refinenet" target="_blank" rel="noopener nofollow ugc">GitHub - DrSleep/light-weight-refinenet: Light-Weight RefineNet for Real-Time...</a></h3>

  <p>Light-Weight RefineNet for Real-Time Semantic Segmentation - GitHub - DrSleep/light-weight-refinenet: Light-Weight RefineNet for Real-Time Semantic Segmentation</p>


  </article>

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

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

<p>The training objective consists of two losses, focal loss for segmentation and scale invariant loss for regression of a distance map.</p>
<p>Training without mixed precision works very well but adding mixed precision makes it diverge.</p>
<p>I tracked the evolution of the scale factor during training and it seems that it just keeps growing and that causes the problem in my case.</p>
<p>In the following I set <code>init_scale = 500</code> and <code>growth_interval=10000</code> which seems to behave better but still suffers from the same divergence problem.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/d/0d82844f5ec608be4bf19df9f28afa93679d202e.png" data-download-href="https://discuss.pytorch.org/uploads/default/0d82844f5ec608be4bf19df9f28afa93679d202e" title="losses2"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/d/0d82844f5ec608be4bf19df9f28afa93679d202e_2_690x371.png" alt="losses2" data-base62-sha1="1VvPXQ1pWa9rkbTacbRc5n6aqyq" width="690" height="371" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/d/0d82844f5ec608be4bf19df9f28afa93679d202e_2_690x371.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/d/0d82844f5ec608be4bf19df9f28afa93679d202e_2_1035x556.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/d/0d82844f5ec608be4bf19df9f28afa93679d202e_2_1380x742.png 2x" data-dominant-color="F9F8F8"><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">losses2</span><span class="informations">1594×859 98.7 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><br>
(l1 loss is not part of the training objective).</p>
<p>Also, I am using torch 1.9.1 with cuda 11.0 and a Titan RTX GPU for this task.</p>

ptrblck: <p>A large scaling factor is beneficial as long as no overflows are seen as it would avoid underflows in the gradient calculation. Once the scaling factor is increased and is indeed causing overfllows, the <code>optimizer.step()</code> method will be skipped in <code>scaler.step</code> and the scale factor decreased again, thus the parameters won’t be updated at all.<br>
I would recommend to update PyTorch to the latest release and check if you are still seeing the issue.</p>

Tariq: <p>Thank you for the reply.</p>
<p>I found that the code base that I was working with had implemented a <code>nan_to_num</code> operation that made the gradscaler unable to work correctly.</p>
<p>The code works fine now.</p>