NAN of batchnorm in AMP

sakura: <p>Hi, we have a model where some layers are</p>
<pre><code class="lang-auto">...
self.shared_1 = nn.Conv2d(512, 64, kernel_size=3, padding=1, bias=False)
self.shared_2 = nn.Batchnorm2d(64, eps=1e-3)
self.shared_3 = nn.ReLU(True)
...
</code></pre>
<p>The forward is</p>
<pre><code class="lang-auto">x1 = self.shared_1(x)
x2 = self.shared_2(x1)
x3 = self.shared_3(x2)
</code></pre>
<p>When we using AMP, we found that the <strong>running_mean</strong> and <strong>runnning_var</strong> of <code>self.shared_2</code> are normal before receiving x1. After producing x2, however, there are some NANs in its  <strong>running_mean</strong> and <strong>runnning_var</strong>. Maybe these values calculated from x1 exceed the range of half precision?</p>
<p>Are there any elegant approach to avoid NAN? Thanks</p>
<p>For instance, how could we keep above conv layer as FP32?</p>

ptrblck: <aside class="quote no-group" data-username="sakura" data-post="1" data-topic="174877">
<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/sakura/48/21826_2.png" class="avatar"> sakura:</div>
<blockquote>
<p>Maybe these values calculated from x1 exceed the range of half precision?</p>
</blockquote>
</aside>
<p>This should not be the case since <code>autocast</code> will use <code>float32</code> in batchnorm layers.<br>
Most likely <code>x1</code> already contained invalid values (Infs or NaNs) and is thus updating the running stats with invalid values, too.<br>
Did you make sure the input does not contain any invalid values?</p>

sakura: <p>We check the min, max, mean values of x1 which are all normal.</p>
<p>Besides, the eps values in all codes are set to 1e-3, as well as the torch.log(). So, does the calculation of <strong>running_mean</strong> and <strong>running_var</strong> in bn is still fp16, which exceeds range?</p>

ptrblck: <p>No, the batchnorm layers will use <code>float32</code> inside the <code>autocast</code> context as already explained. Unless you are calling <code>half()</code> on this model and use a manual <code>flaot16</code> training, this operation should not overflow.</p>

sakura: <p>We know that the <code>autocast(enable=False)</code> could be used in <code>forward()</code> to set specific layers as FP32. Are there some approaches to keep specific layers as FP32 in <code>__init__</code> function?</p>
<p>Maybe we could find bug layer by progressively replacing FP16 layers as FP32…</p>

vadimkantorov: <p><a class="mention" href="/u/sakura">@sakura</a> Please support this simple autocast-module-wrappers feature request at <a href="https://github.com/pytorch/pytorch/issues/70166" class="inline-onebox" rel="noopener nofollow ugc">[feature request] Autocast module and function wrappers · Issue #70166 · pytorch/pytorch · GitHub</a>, so far it hasn’t obtained much support/traction <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<p>Probably it could be as simple as:</p>
<pre><code class="lang-python">class AutocastModule(torch.nn.Module):
  def __init__(self, module, **autocast_options):
    super().__init__()
    self.module = module
    self.autocast_options = autocast_options

  def forward(self, *args, **kwargs):
    with torch.autocast(**self.autocast_options):
      return self.module(*args, **kwargs)
</code></pre>

sakura: <p>Thanks. The bug has been solved : ）</p>

Zane: <p>I got same bug on my model,how did you solve this problem?<br>
When training, every think is normal, but when it comes to val, the running_mean and var become nan</p>

vadimkantorov: <p>One solution is to do manual updates of running stats in separate, manually managed tensors. Like this, you can filter out the updates/bad values that you don’t want.</p>
<p>One can do it by setting <code>bn.momentum = 1</code>, and then <code>bn.running_mean</code>/<code>bn.running_var</code> would actually contain batch statistics. Then you would use these batch statistics to update the running statistics tensors that you keep on the side. Before evaluation, you need to copy running statistics into <code>bn.running_mean</code>/<code>bn.running_var</code></p>