How does batchNorm2d work? (with CoLab code)

JH_S: <p>Hi,</p>
<p>I organize my problem below.<br>
I hope you utilize <a href="https://colab.research.google.com/drive/12Khe2duj4OeNfHBtvTKd3g9edfA8b--J?usp=sharing" rel="noopener nofollow ugc">this colab code</a> I made.</p>
<h1><a name="h-1-implement-new-batchnorm-class-1" class="anchor" href="#h-1-implement-new-batchnorm-class-1"></a>1. Implement new Batchnorm class</h1>
<p>I implemented a new BatchNorm2d that is mathematically identical to the original one but has a different coding style.</p>
<pre data-code-wrap="python"><code class="lang-python">class BatchNorm2d_new(nn.BatchNorm2d):
    def forward(self, input):
        self._check_input_dim(input)

        statistics_tbn = {'mean':input.mean(dim=[0, 2, 3]).detach().clone(), 'var':input.var(dim=[0, 2, 3]).detach().clone()}
        statistics_tbn['inv_std'] = torch.reciprocal(torch.sqrt((statistics_tbn['var']+self.eps)))

        output = input -  _unsqueeze_(statistics_tbn['mean'])
        output = output * _unsqueeze_(statistics_tbn['inv_std'])
        output = output * _unsqueeze_(self.weight)
        output = output + _unsqueeze_(self.bias)
        return output

def _unsqueeze_(tensor):
    return tensor.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)

# This function change origin BN to New BN in our network, reclusively.
def convert_bn(module):
    module_output = module
    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
        module_output = BatchNorm2d_new(
            module.num_features,
            module.eps,
            module.momentum,
            module.affine
        )
        module_output.weight.data = module.weight.data.clone()
        module_output.bias.data = module.bias.data.clone()
        module_output.running_mean = module.running_mean
        module_output.running_var = module.running_var
              
    for name, child in module.named_children():
        module_output.add_module(
            name, convert_bn(child),
        )
    del module
    return module_output
</code></pre>
<h1><a name="h-2-train-on-cifar10-with-origin-bn-2" class="anchor" href="#h-2-train-on-cifar10-with-origin-bn-2"></a>2. Train on CIFAR10 with origin BN</h1>
<pre data-code-wrap="python"><code class="lang-python">from torchvision.models import resnet18
net = resnet18(pretrained=True)

# Train Resnet18 on CIFAR10
# Results
# accuracy on train dataset (not eval dataset)
[1,     1] loss: 0.007 acc: 0.000
[1,    41] loss: 0.094 acc: 0.500
[1,    81] loss: 0.035 acc: 0.438
[1,   121] loss: 0.029 acc: 0.484
[1,   161] loss: 0.027 acc: 0.484
[1,   201] loss: 0.025 acc: 0.641
[1,   241] loss: 0.024 acc: 0.594
[1,   281] loss: 0.023 acc: 0.719
</code></pre>
<h1><a name="h-3-train-on-cifar10-with-new-bn-3" class="anchor" href="#h-3-train-on-cifar10-with-new-bn-3"></a>3. Train on CIFAR10 with new BN.</h1>
<pre data-code-wrap="python"><code class="lang-python">from torchvision.models import resnet18
net = resnet18(pretrained=True)
net = convert_bn(net)

# Train Resnet18 on CIFAR10
# Results
# accuracy on train dataset (not eval dataset)
[1,     1] loss: 0.007 acc: 0.000
[1,    41] loss: 0.158 acc: 0.094
[1,    81] loss: 0.052 acc: 0.141
[1,   121] loss: 0.046 acc: 0.172
[1,   161] loss: 0.046 acc: 0.125
[1,   201] loss: 0.045 acc: 0.266
[1,   241] loss: 0.046 acc: 0.094
</code></pre>
<h1><a name="h-4-question-4" class="anchor" href="#h-4-question-4"></a>4. Question</h1>
<p>As you can see, the accuracy with new BN is too bad then one with original BN. In my opinion, Two BN class should behave the same, since they are mathematically identical. (following <a href="https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html" rel="noopener nofollow ugc">torch.batchnorm</a>)<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/0/5/0540e5ab6c43aa3ed27b5bb4e7ac304781bc4fba.png" alt="image" data-base62-sha1="KtqAuzfGEPL0xyLcsderbLGYMG" width="226" height="60"><br>
Q: <strong>Why isn’t learning done with new BN? How should I modify the code?</strong></p>
<p>I hope I explain my question clearly.<br>
Best regards <img src="https://discuss.pytorch.org/images/emoji/apple/smiling_face.png?v=12" title=":smiling_face:" class="emoji" alt=":smiling_face:" loading="lazy" width="20" height="20"></p>
<h2><a name="ps-5" class="anchor" href="#ps-5"></a>PS.</h2>
<ol>
<li>They trained on train mode (net.train()). I think this is not related to Train &amp; Eval mode. This problem is all about weight(gamma) and bias(beta).</li>
<li>Calculating loss and accuracy are implemented in <a href="https://colab.research.google.com/drive/12Khe2duj4OeNfHBtvTKd3g9edfA8b--J?usp=sharing" rel="noopener nofollow ugc">colab code</a>.</li>
</ol>

JH_S: <p>I solve a problem after changing my code.</p>
<h2>
<a name="from-1" class="anchor" href="#from-1"></a>From</h2>
<p>In BatchNorm2d_new class,</p>
<pre><code class="lang-python">statistics_tbn = {'mean':input.mean(dim=[0, 2, 3]).detach().clone(), 'var':input.var(dim=[0, 2, 3]).detach().clone()}
statistics_tbn['inv_std'] = torch.reciprocal(torch.sqrt((statistics_tbn['var']+self.eps)))
</code></pre>
<h2>
<a name="to-2" class="anchor" href="#to-2"></a>To</h2>
<pre><code class="lang-python">mean = x.mean(dim=[0, 2, 3], keepdim=True)
var = x.var(dim=[0, 2, 3], keepdim=True)
</code></pre>
<p><strong>It was because of <code>detach</code>. But I still don’t know why <code>detach</code> results in non-training.</strong></p>

azhanmohammed: <aside class="quote no-group" data-username="JH_S" data-post="2" data-topic="161614">
<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/jh_s/48/40536_2.png" class="avatar"> JH_S:</div>
<blockquote>
<p><strong>It was because of <code>detach</code>. But I still don’t know why <code>detach</code> results in non-training.</strong></p>
</blockquote>
</aside>
<p>Using <code>detach</code> removes the layer or variable from the model’s training graph, so while the losses are being back propagated the detached outputs do not contribute to the loss, and similarly they are not updated as well. So using <code>detach</code> would result in non-training of the layers.</p>