Batch Norm Broadcasting

gabrieldernbach: <p>When working with vectorial data, I sometimes need to leave the <code>batch x dimension</code> format in favour of <code>batch x samples x dimension</code>.</p>
<p>The sampling dimension can originate from the latent space of importance weighted autoencoders, or just multiple measurements of the same instance during an experiment.</p>
<p>While linear layers easily broadcast along the additional dimensions, batch norm layers need special care here.</p>
<p>My work around about the missing bach norm broadcasting is a double reshape<br>
from <code>reshape(batch * samples x dimension) =&gt; batch_norm(...) =&gt; reshape(batch, samples, dimension</code></p>
<p>Is there a way to extend batch norm without the supposedly inefficient reshaping?</p>
<p>For illustration</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

lin = nn.Linear(10, 10)

b, p, d = 16, 80, 10
x = torch.randn(b, p, d)
x_ = x.reshape(b * p, d)

o = lin(x)
o_ = lin(x_)

print(torch.allclose(o, o_.reshape(b, p, d)))
"&gt;&gt; yay :)"


bn = nn.BatchNorm1d(10)
bn(x_) # is ok
bn(x)
"&gt;&gt; nay :?"
"RuntimeError: running_mean should contain 80 elements not 10"</code></pre>

ptrblck: <p>Batchnorm layers use running stats internally and expect the input channels to match their dimension.<br>
The <code>x</code> input is using <code>80</code> channels (<code>dim1</code>), which raises the error and I’m unsure which workaround should be used instead.<br>
Could you explain your idea a bit more?</p>

gabrieldernbach: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> thanks for pointing out the <code>channel</code> dimension, I revisited the docstring. I only had vectorial data (<code>N x L</code> input size) in mind, and did not think about the use case of convolution with (<code>N x C x L</code> input size).</p>
<p>The current behaviour of batch norm on tensors (<code>N x C x L</code>) is designed to conform with what should be used with 1D convolutional outputs. That is, each channel will get its own <code>mean</code> and <code>var</code>, right? And all elements in the channel will be standardised with the same <code>mean</code>, <code>var</code> pair.</p>
<p>I can then reframe my issue:<br>
The default behaviour of BatchNorm on tensors (instead of matrices), does not broadcast the matrix operation (element wise batch norm), but completely switches to a different kind operation (channel wise batch norm).</p>
<p>I appreciate the fact, that this may be a sound default for the many users working with CNNs (including me <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> ). I however see many cases, where an actual broadcasting of the BatchNorm layer to higher order tensors is desired. The broadcasting would then be identical to how linear layers cope with higher dimensional tensors (only apply to last dimension, broadcast else).</p>
<p>Maybe this is too didactic, but it could be helpful, if the library explicates the difference of applying an elementwise batch normalization to a channel wise batch normalization by setting a default flag.</p>

ptrblck: <p>Based on your description it seems you might be looking for <a href="https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html#torch.nn.LayerNorm"><code>LayerNorm</code></a>?<br>
From the docs:</p>
<blockquote>
<p>Unlike Batch Normalization and Instance Normalization, which applies scalar scale and bias for each entire channel/plane with the <code>affine</code> option, Layer Normalization applies per-element scale and bias with <code>elementwise_affine</code> .</p>
</blockquote>