Correct function for BatchNorm2D

sonnguyen: <p>I’m trying to get the correct function of batchnorm2d.<br>
Suppose my model is the following</p>
<blockquote>
<p>class Net(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(Net, self).<strong>init</strong>()<br>
self.conv1 = nn.Conv2d(1, 32, 5,padding=2)<br>
self.pool = nn.MaxPool2d(2, 2)<br>
self.bn1 = torch.nn.BatchNorm2d(32)<br>
self.fc1 = nn.Linear(32<em>14</em>14, 10)</p>
<pre><code>def forward(self, x):
    x = self.bn1(self.pool(self.conv1(x))) 
    x = F.relu(x)
    x = x.view(-1, 32 * 14 * 14)
    x = self.fc1(x)
    return x
</code></pre>
<p>net=Net()<br>
tem1 = net.pool(net.conv1(x))</p>
</blockquote>
<p>Could you please tell me what is the equivalent of below equation?</p>
<blockquote>
<p>bn = net.bn1(tem1)</p>
</blockquote>
<p>The most closed equation I can have is</p>
<blockquote>
<p>net.bn1.weight.view([32,1,1])*((tem1 - net.bn1.running_mean.view([32,1,1]))/((net.bn1.running_var.view([32,1,1]) + net.bn1.eps).sqrt())) + net.bn1.bias.view([32,1,1])</p>
</blockquote>
<p>But still it has a slight difference with net.bn1(tem1)<br>
Many thanks.</p>

Kushaj: <p>What do you mean by equivalent? As far as I understand, you want to simplify</p>
<pre><code class="lang-python">tem1 = net.pool(net.conv1(x))
bn = net.bn1(tem1)
</code></pre>
<p>Which is just a combination of conv-&gt;pool-&gt;batchnorm</p>

sonnguyen: <p>I want to understand what is the exact implementation inside the batchnorm2d. So I’m trying to figure out the exact equation inside the net.bn1(tem1). Equivalent means they should have the same output.</p>

InnovArul: <p>I think, <code>running_mean</code> and <code>running_var</code> are computed during training, however they are used for normalization only during testing. During training, the batch statistics (<code>mean</code> and <code>std</code>) are directly used for normalization. <a class="mention" href="/u/ptrblck">@ptrblck</a> Is that right?.</p>
<p>As far as <code>BatchNorm2d</code> computation is concerned, the following works:</p>
<pre><code class="lang-auto">bn_in = ... # input tensor for batch norm

# squeeze spatial dimensions
bn_in_flat = bn_in.view(bn_in.shape[0], bn_in.shape[1], -1)

# we need mean and std per channel across whole batch. so permute the batch dimension and flatten
bn_in_flat = bn_in_flat.permute(1,0,2).view(bn_in.shape[1], -1)

# calculate mean
bn_in_mean = bn_in_flat.mean(dim=-1).view(1, bn_in.shape[1], 1,1)

# calculate unbiased std
bn_in_std = bn_in_flat.std(dim=-1, unbiased=False).view(1, bn_in.shape[1], 1,1)

# batch norm 2D computation
bn_out = ((net.bn1.weight.view(1, out_channels, 1, 1) * 
            (bn_in - bn_in_mean) / (bn_in_std**2 + net.bn1.eps).sqrt()) + 
         net.bn1.bias.view(1,out_channels,1,1))
</code></pre>

ptrblck: <p>Yes, you are correct. The running estimates are used during <code>eval</code>. If <code>track_running_stats</code> is set to <code>False</code>, the batch statistics will be used during training and eval.</p>

InnovArul: <aside class="quote no-group quote-modified" data-username="ptrblck" data-post="5" data-topic="41045">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>If <code>track_running_stats</code> is set to <code>True</code> , the batch statistics will be used during training and eval.</p>
</blockquote>
</aside>
<p>A small correction/clarification is needed. <code>track_running_stats</code> is set to <code>False</code>?</p>

ptrblck: <p>Yes, thanks for catching this typo! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=7" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

sonnguyen: <p>Hi all, thank much for your help. I’m able to generate the same output now.</p>

Taro_Kiritani: <p>I am wondering why <code>running_mean</code> and <code>running_var</code> are not used during training. Intuitively, that would make the model more stable, and the convergence would be faster.</p>

ptrblck: <p>These running stats are updated during training and are thus set to the default values of <code>0s</code> for the <code>running_mean</code> and <code>1s</code> for the <code>running_var</code>. Using them during training without updating wouldn’t normalize the data.</p>
<p>Also, even if you update the stats during training but use them as well, I would assume your model might yield quite a bad performance, since the underlying mean and var might differ from the default values and the model would be trained on “unnormalized” data at the beginning.<br>
Batchnorm layers deeper in the model would also get activation values, which were not created using normalized data, so your complete model could also break.</p>
<p>EDIT: just a quick update:<br>
as you see there are a lot of "could/should"s in my post, so don’t let my post stop you from experimenting with this approach, if you think it can work fine. Also, please update me once you’ve run some experiments. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

Taro_Kiritani: <p>After some research, I found this approach (a bit modified version) has been used in <a href="https://arxiv.org/abs/1702.03275" rel="noopener nofollow ugc">batch renorm</a>. It is not included in pytorch yet, but some third party repos like <a href="https://github.com/ludvb/batchrenorm" rel="noopener nofollow ugc">this one</a> are available.</p>

sigma_x: <p>So is it correct that at test/eval time, only a single mean/std used, not the whole history that <code>running_mean</code> and <code>running_var</code> store?</p>

ptrblck: <p>The running stats are not storing a history of the values, but are updated to a single value as described in the <a href="https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html#torch.nn.BatchNorm2d">docs</a> using the current batch stats, the running stats, and the <code>momentum</code>.</p>

sigma_x: <p>OK, so there are 64 values in <code>running_mean</code> and <code>running_var</code>. Which ones are used at test time, with <code>model.eval()</code>?</p>

ptrblck: <p>All of these values will be used, since the <code>running_mean</code> and <code>running_var</code> contain <code>num_features</code> values, which corresponds to the number of input channels.</p>

sigma_x: <p>OK, I think I got it: during training only the current per-channel mean/std are used (+scale and shift factors), at test time the history (also per-channel).</p>