Getting parameters of `torch.nn.BatchNorm2d` during training

username: nullgeppetto
cooked: <p>Hi all, I’m trying to figure out how exactly a 2d BN layer applies on an input tensor. Let me give a toy example:</p>
<pre><code class="lang-auto">import torch                                                                                                                                                                          
import torch.nn as nn 

num_features = 3                                                                                                                                                                      
X = torch.randn(2, num_features, 4, 4) 
bn = nn.BatchNorm2d(num_features)

Y = bn(X)
</code></pre>
<p>What I would like to do is to “reproduce” the output of the BN layer, <code>Y</code>, not by passing <code>X</code> through <code>bn()</code>, but by using the formula shown <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.BatchNorm2d" rel="nofollow noopener">here</a>.</p>
<p>That is, I would something like:</p>
<pre><code class="lang-auto">Z = bn.xxx * torch.div(X - bn.yyy, torch.sqrt(bn.zzz + bn.eps)) + bn.ppp
</code></pre>
<p>I’m looking for the correct attributes of <code>bn</code> (<code>xxx</code>, <code>yyy</code>, <code>zzz</code>, and <code>ppp</code>) - or, even a yet better way to do that, since -during training- I want to get the parameters of BN layer and apply them to another tensor.</p>
<p>Thanks a lot.</p>

username: sami
cooked: <p>If I understood correctly, you want to incorporate the batch norm parameters into e.g. a convolution,  then you can do e.g. this:</p>
<pre><code class="lang-auto">    def incorporateBatchNorm(self, bn):

        gamma = bn.weight
        beta = bn.bias
        mean = bn.running_mean
        var = bn.running_var
        eps = bn.eps

        var_sqrt = torch.sqrt(var + eps)

        w = (self.weight * gamma.reshape(self.out_channels, 1, 1, 1)) / var_sqrt.reshape(self.out_channels, 1,
                                                                                         1, 1)
        b = ((self.bias - mean) * gamma) / var_sqrt + beta

        self.weight = nn.Parameter(w)
        self.bias = nn.Parameter(b)
</code></pre>
<p>And then you can just do the forward pass of the convolution layer.</p>

username: nullgeppetto
cooked: <p>Hi <a class="mention" href="/u/sami">@sami</a>, thank you very much for your answer. I understand what you did above, and even though I don’t want to use it in a convolutional (or any other parametric operation), your approach should do the job for me. But it doen’t…</p>
<p>Let’s say that you have a 4-dimensional tensor <code>X</code> of size <code>(bs, ch, d, d)</code>:</p>
<pre><code class="lang-auto">bs = 2, ch=3, d=4
X = torch.randn(bs, ch, d, d)
</code></pre>
<p>Then, <code>Y = bn(X)</code> will forward <code>X</code> through the BN layer, by computing <code>bn.running_mean</code>, <code>bn.running_var</code>, etc…</p>
<p>I would like to reproduce the value of <code>Y</code> by manually applying the operation of BN. That should be the following:</p>
<pre><code class="lang-auto">Y_ = bn.weight.reshape(1, ch, 1, 1) * (X - bn.running_mean.weight.reshape(1, ch, 1, 1))/torch.sqrt(bn.running_var.weight.reshape(1, ch, 1, 1) + bn.eps) + bn.bias.weight.reshape(1, ch, 1, 1)
</code></pre>
<p>But it doesn’t seem to work. I tried also with a BN layer without any affine transformation, but it still cannot reproduce <code>Y</code>.</p>

username: ZacharyGong
cooked: <p>Hi, <a class="mention" href="/u/nullgeppetto">@nullgeppetto</a>, I’m facing the same problem as yours.  I wonder if you have found the answer?</p>

username: nullgeppetto
cooked: <p>Hi <a class="mention" href="/u/zacharygong">@ZacharyGong</a>, the truth is that abandoned the idea quickly so I didn’t try to figure it out. Although I would still want to sort it out. I may take another look soon and let you know, but I think that we both miss something here (it’s pretty much straightforward after all). Maybe <a class="mention" href="/u/sami">@sami</a> wants to take a look also (sorry for the spam!).</p>

username: nemokwy
cooked: <blockquote>
<p>t = image<br>
t = F.conv2d(t, network.conv1.weight)<br>
w = torch.ones(t.size()) * network.bn1.weight.reshape(1,-1,1,1)<br>
b = torch.ones(t.size()) * network.bn1.bias.reshape(1,-1,1,1)<br>
m = torch.ones(t.size()) * network.bn1.running_mean.reshape(1,-1,1,1)<br>
v = torch.ones(t.size()) * network.bn1.running_var.reshape(1,-1,1,1)<br>
v = v.sqrt()<br>
t = (t - m) / v * w + b<br>
t[0][0][0]</p>
</blockquote>
<p>The code above got exact the same results as</p>
<blockquote>
<p>self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, bias=False)<br>
self.bn1 = nn.BatchNorm2d(6)</p>
</blockquote>
<blockquote>
<p>def forward(self, t):<br>
t = self.bn1(self.conv1(t))</p>
</blockquote>
<p>in network.eval() mode. <a class="mention" href="/u/zacharygong">@ZacharyGong</a> <a class="mention" href="/u/nullgeppetto">@nullgeppetto</a> Hope it helps.</p>