nn.BatchNorm vs MyBatchNorm

Neofytos: <p>I have reimplemented BatchNorm1D based on the implementation provided by <a class="mention" href="/u/ptrblck">@ptrblck</a> (greatly appreciated!), here: <a href="https://github.com/ptrblck/pytorch_misc/blob/master/batch_norm_manual.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch_misc/batch_norm_manual.py at master · ptrblck/pytorch_misc · GitHub</a></p>
<p>In order to verify identical behaviour with the nn.BatchNorm equivalent, I initiate 2 models (as well as 2 optimizers), one using MyBatchNorm and one using nn.BatchNorm. Although I can verify that in the beginning, “compareBN” returns identical results (<code>Max diff:  tensor(8.5831e-06, device='cuda:0', grad_fn=&lt;MaxBackward1&gt;)</code>), as the training progresses, very tiny differences begin to emerge.</p>
<p>A few iterations in:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/9/19c3014e50a882bd685ac333cbbd4abac5bf5c78.png" data-download-href="https://discuss.pytorch.org/uploads/default/19c3014e50a882bd685ac333cbbd4abac5bf5c78" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/19c3014e50a882bd685ac333cbbd4abac5bf5c78_2_690x232.png" alt="image" data-base62-sha1="3FTJsFi26MMw3qAImsfIKwmA8vm" width="690" height="232" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/19c3014e50a882bd685ac333cbbd4abac5bf5c78_2_690x232.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/19c3014e50a882bd685ac333cbbd4abac5bf5c78_2_1035x348.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/9/19c3014e50a882bd685ac333cbbd4abac5bf5c78_2_1380x464.png 2x" data-dominant-color="D3D4CA"><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">image</span><span class="informations">1659×558 95.6 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>More iteratations later:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/0/30ec5c6f8dfce1a9990ebe7f76badbc6015de8ca.png" data-download-href="https://discuss.pytorch.org/uploads/default/30ec5c6f8dfce1a9990ebe7f76badbc6015de8ca" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/0/30ec5c6f8dfce1a9990ebe7f76badbc6015de8ca_2_690x196.png" alt="image" data-base62-sha1="6YNjPw9cahXbmoXfYex3BoBERFU" width="690" height="196" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/0/30ec5c6f8dfce1a9990ebe7f76badbc6015de8ca_2_690x196.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/0/30ec5c6f8dfce1a9990ebe7f76badbc6015de8ca_2_1035x294.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/0/30ec5c6f8dfce1a9990ebe7f76badbc6015de8ca_2_1380x392.png 2x" data-dominant-color="C7C8BB"><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">image</span><span class="informations">1657×473 71.2 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>By the end of the 25 epochs, both the models and the results differ significantly. Note that I have set all random seeds to 0 (including deterministic behaviour for cudnn and cuda).</p>
<p>This reminded me of fp16 vs fp32 training. Can it be the case that the underlying implementation of nn.BatchNorm uses different precision? Or is it something else? The MyBatchNorm method is also slightly slower (+2 seconds on a 10 second epoch), but I’d assume that this is due to C++ being used for nn.BatchNorm. I used a 2 layer FCNN with 2 BatchNorm layers on MNIST. Interestingly, the 2nd BatchNorm layer is slower in showing disrepancies.</p>

Neofytos: <p>I was able to get the models (predictions as well as the batchnorm params) closer to each other by using double precision (<code>model.double()</code>). In fact, their differences now have been stabilised around 10^-10 - 10^-11 (based on the output of <code>print('Max diff: ', (outputs - outputs_).abs().max()</code> and <code>compareBN</code>), whereas before it would at some point escalate to 0.5-1.5 and therefore end up with 2 completely different models.</p>
<p>The downside is that it now takes longer to train due to the double precision. If anyone has any recommendation as to how to ameliorate the problem without the overhead, I am all ears <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ptrblck: <p>My manual implementation was originally intended to be used as a reference implementation to see, how the running estimates are updated and how the batchnorm layer works in general.<br>
The performance difference is expected, as I used native methods without specific (and faster) kernels.</p>
<p>The difference in after a couple of iterations is most likely created due to the accumulation of rounding errors created by the limited floating point precision.<br>
It’s also an indication for this, since you see a reduction in these errors using <code>float64</code>.</p>
<p>What’s your use case, that you would like to use this implementation in a real model? <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

Neofytos: <p>Could you be more specific with what you mean with faster kernels :)? I’d be interested to look into that. I am trying to implement ghost batch normalization (in essence, estimate batch statistics based on smaller sample sizes than the original mini-batch) for research purposes.</p>

ptrblck: <p>I’m not calling into any batchnorm functions defined in <a href="https://github.com/pytorch/pytorch/blob/d035d05080729c30636ff30fcc068de3c7e9badd/aten/src/ATen/native/Normalization.cpp">Normalization.cpp</a>, such as <a href="https://github.com/pytorch/pytorch/blob/d035d05080729c30636ff30fcc068de3c7e9badd/aten/src/ATen/native/Normalization.cpp#L426-L432">the cudnn implementation</a>, but use just native operators.</p>

Neofytos: <p>Other than the speedup one would gain from adopting those functions, do these functions also facilitate determinism? When I don’t use cuda(), I can reproduce the experiments, however, when I use cuda() and my own implementation of BatchNorm, it is not reproducible <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>

ptrblck: <p>Have a look at the <a href="https://pytorch.org/docs/stable/notes/randomness.html">Reproducibility docs</a>, if not already done.<br>
What exactly did you change in the implementation, as I’m able to get deterministic outputs just be seeing the code with my code snippet.</p>

Neofytos: <p>Hey, thanks for following-up. I appreciate it.</p>
<blockquote>
<p>Have a look at the <a href="https://pytorch.org/docs/stable/notes/randomness.html" rel="noopener nofollow ugc">Reproducibility docs</a>, if not already done.</p>
</blockquote>
<p>The following is included with every run:</p>
<pre><code class="lang-auto">torch.manual_seed(0)
torch.cuda.manual_seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(0)
random.seed(0)
os.environ['PYTHONHASHSEED'] = str(0)
</code></pre>
<blockquote>
<p>What exactly did you change in the implementation, as I’m able to get deterministic outputs just be seeing the code with my code snippet.</p>
</blockquote>
<p>You are right, I modified the implementation back to yours, and it’s reproducible. It’s very strange because the nondeterminism seems to arise from rounding errors that do not happen in the same way, i.e. on 2 different runs I get:<br>
Epoch 1/2, training loss: 0.2351639289064084, validation loss: 0.09888318485801263<br>
Epoch 1/2, training loss: 0.23525172622356727, validation loss: 0.10040923198367216</p>
<p>I’ve added the following: reshape input tensor, take the mean &amp; var across the newly introduced dimension, and then use repeat_interleave on mean &amp; var. Finally, do the (input - mean) / sqrt(var) calculation.</p>

ptrblck: <p>Ah OK.<br>
The backward pass of <code>repeat_interleave</code> is not deterministic as explained in the linked docs:</p>
<blockquote>
<p>Additionally, the backward path for  <code>repeat_interleave()</code>  operates nondeterministically on the CUDA backend because  <code>repeat_interleave()</code>  is implemented using  <code>index_select()</code> , the backward path for which is implemented using  <code>index_add_()</code> , which is known to operate nondeterministically (in the forward direction) on the CUDA backend (see above).</p>
</blockquote>

loganfrank: <p>I’m also trying to reimplement BatchNorm2d based off your provided approach and getting similar rounding errors. Do you have any recommendations for a way to completely remove these differences?</p>
<p>My use case is I’m trying to separate the Batch Norm layer into two separate parts - mean/var (“head”) and weight/bias (“tail”). Ideally, I would like to have each part as a separate class that extends nn.Module and <code>self.tail(self.head(x))</code> would be the exact same as doing <code>self.batch_norm(x)</code>. You could also think of the “head” as doing <code>nn.BatchNorm2d(num_features, affine=False)</code>, but I just need a way to do the affine transformation without any precision errors from the normal method.</p>

ptrblck: <aside class="quote no-group" data-username="loganfrank" data-post="10" data-topic="82682">
<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/loganfrank/48/35983_2.png" class="avatar"> loganfrank:</div>
<blockquote>
<p>but I just need a way to do the affine transformation without any precision errors from the normal method.</p>
</blockquote>
</aside>
<p>If you are seeing absolute errors in the range <code>1e-5</code> they would be created by the limited floating point precision and the order of operations.<br>
I don’t think you can easily guarantee the order of floating point operations in the manual approach to match the rounding errors in the internal PyTorch version.<br>
If you need more precision, you could use <code>.double()</code> types for the data, buffers, and parameters.</p>