Why 2D batch normalisation is used in features and 1D in classifiers?

4deep.prk: <p>What is the difference between BatchNorm2d and BatchNorm1d? Why a BatchNorm2d is used in features and BatchNorm1d is used in classifier?</p>

Nikronic: <p>Hi,</p>
<p>There is no mathematical difference between them, except the dimension of input data.<br>
<a href="https://pytorch.org/docs/master/generated/torch.nn.BatchNorm2d.html" rel="nofollow noopener">nn.BatchNorm2d</a> only accepts 4D inputs while <a href="https://pytorch.org/docs/master/generated/torch.nn.BatchNorm1d.html" rel="nofollow noopener">nn.BatchNorm1d</a> accepts 2D or 3D inputs. And because of that, in <code>features</code> which has been constructed of <code>nn.Conv2d</code> layers, inputs are <code>[batch, ch, h, w]</code> (4D) we need <code>BatchNorm2d</code> and in <code>classifier</code> we have <code>Linear</code> layers which accept <code>[batch, length]</code> or <code>[batch, channel, length]</code> (2D/3D) so we need <code>BatchNorm1d</code>.<br>
Two linked docs completely explain this idea.</p>
<p>Bests</p>

level1807: <p>If there is no difference between them, then why would there be two different functions? And why wouldn’t there be a universal BatchNorm class that accepts inputs with arbitrary dimensions?</p>

Nikronic: <p>Hi</p>
<p>There is a universal BatchNorm!<br>
Simply put here is the architecture (<a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/batchnorm.html#BatchNorm2d" rel="noopener nofollow ugc">torch.nn.modules.batchnorm — PyTorch 1.11.0 documentation</a>):</p>
<ol>
<li>a base class for <em>normalization</em>, either <code>Instance</code> or <code>Batch</code> normalization → <code>class _NormBase(Module)</code>. This class includes no computation and does not implement <code>def _check_input_dim(self, input)</code></li>
<li>Now we have <code>class _BatchNorm(_NormBase)</code> that extends <code>_NormBase</code> which actually does the computation and tracks the values necessary for it. In the last line, you see that class calls <code>return F.batch_norm(..)</code> in its <code>forward</code> function. We will talk about <code>F.batch_norm()</code> in a bit.</li>
<li>In the end, you have classes in form <code>class BatchNormXd(_BatchNorm)</code> that extend <code>_BatchNorm</code> and the only thing they do is to implement <code>_check_input_dim(self, input)</code> that was intentionally left behind in <code>_NormBase</code> class (see step 1).</li>
</ol>
<p>About <code>torch.nn.functional.batch_norm</code> (from step 2): This function, does all the computation given values, in other words, if you want to mix all parts into a single class, you should be able too, but it would break the modularity, etc.</p>
<p>Here is an example that you can exactly replicate what <code>BatchNormXd</code> does for just a single forward pass given the normalization formulation:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/e/4e8d2a4d3700d46ab091a43d591e2e7664beb9d8.png" alt="image" data-base62-sha1="bcTGE0Lb9Ld0Z7cuGUBuoSTYEXu" width="248" height="74"></p>
<pre><code class="lang-auto"># Test case
x = torch.randn(2, 7)  # batch=2, features=7
running_mean = x.mean(dim=0)  # assuming 'mean' tracked during training
# Remark: as the documentation says, we must use biased estimator, i.e. 'unbiased=False'.
running_var = x.var(dim=0, unbiased=False)  # assuming 'var' tracked during training
gamma = None  # assuming 'gamma' is not set
beta = None  # assuming 'beta' is not set

&gt;&gt;&gt; x
tensor([[ 1.6080,  1.5907, -1.0321,  1.0416, -0.8388,  0.0759, -0.9885],
        [-0.1404,  0.7668,  1.4246, -0.4341, -1.0590,  0.7760,  0.8207]])
		
# BatchNorm as a function
import torch.nn.functional as F

&gt;&gt;&gt; F.batch_norm(x, running_mean, running_var, gamma, beta, momentum=0.)
tensor([[ 1.0000,  1.0000, -1.0000,  1.0000,  0.9996, -1.0000, -1.0000],
        [-1.0000, -1.0000,  1.0000, -1.0000, -0.9996,  1.0000,  1.0000]])


# BatchNorm as a class
import torch.nn as nn

bn1d = nn.BatchNorm1d(x.shape[1], affine=False, momentum=None)  # 'affine=False' sets beta and gamma to None
# you can verify mean and var by bn1d.running_mean and bn1d.running_var
# you can verify gamma and beta by bn1d.weight and bn1d.bias

&gt;&gt;&gt; bn1d(x)
tensor([[ 1.0000,  1.0000, -1.0000,  1.0000,  0.9996, -1.0000, -1.0000],
        [-1.0000, -1.0000,  1.0000, -1.0000, -0.9996,  1.0000,  1.0000]])
</code></pre>
<p>Bests</p>

level1807: <p>that’s very good, thanks for the pointer!</p>