PyTorch BatchNorm2d Calculation
I am trying to understand the mechanics of PyTorch BatchNorm2d through calculation. My example code:
import torch
from torch import nn

torch.manual_seed(123)

a = torch.rand(3,2,3,3)
print(a)

print(nn.BatchNorm2d(2)(a))
#print(a[:,0,:,:])
mean_by_plane_feature = torch.mean(a,dim=0)
std_by_plane_feature = torch.std(a,dim=0)
print(mean_by_plane_feature)
print(std_by_plane_feature)

Output:
tensor([[[[0.2961, 0.5166, 0.2517],
          [0.6886, 0.0740, 0.8665],
          [0.1366, 0.1025, 0.1841]],

         [[0.7264, 0.3153, 0.6871],
          [0.0756, 0.1966, 0.3164],
          [0.4017, 0.1186, 0.8274]]],


        [[[0.3821, 0.6605, 0.8536],
          [0.5932, 0.6367, 0.9826],
          [0.2745, 0.6584, 0.2775]],

         [[0.8573, 0.8993, 0.0390],
          [0.9268, 0.7388, 0.7179],
          [0.7058, 0.9156, 0.4340]]],


        [[[0.0772, 0.3565, 0.1479],
          [0.5331, 0.4066, 0.2318],
          [0.4545, 0.9737, 0.4606]],

         [[0.5159, 0.4220, 0.5786],
          [0.9455, 0.8057, 0.6775],
          [0.6087, 0.6179, 0.6932]]]])
tensor([[[[-0.5621,  0.2574, -0.7273],
          [ 0.8968, -1.3879,  1.5584],
          [-1.1552, -1.2819, -0.9787]],

         [[ 0.5369, -1.0117,  0.3888],
          [-1.9141, -1.4584, -1.0073],
          [-0.6859, -1.7524,  0.9171]]],


        [[[-0.2425,  0.7925,  1.5103],
          [ 0.5422,  0.7042,  1.9901],
          [-0.6425,  0.7846, -0.6311]],

         [[ 1.0298,  1.1880, -2.0520],
          [ 1.2915,  0.5833,  0.5047],
          [ 0.4593,  1.2495, -0.5645]]],


        [[[-1.3761, -0.3375, -1.1132],
          [ 0.3187, -0.1512, -0.8011],
          [ 0.0269,  1.9569,  0.0493]],

         [[-0.2561, -0.6096, -0.0199],
          [ 1.3619,  0.8356,  0.3525],
          [ 0.0933,  0.1281,  0.4116]]]], grad_fn=&lt;NativeBatchNormBackward&gt;)
tensor([[[0.2518, 0.5112, 0.4177],
         [0.6049, 0.3724, 0.6937],
         [0.2885, 0.5782, 0.3074]],

        [[0.6999, 0.5455, 0.4349],
         [0.6493, 0.5804, 0.5706],
         [0.5721, 0.5507, 0.6515]]])
tensor([[[0.1572, 0.1521, 0.3810],
         [0.0784, 0.2829, 0.4042],
         [0.1594, 0.4411, 0.1406]],

        [[0.1723, 0.3110, 0.3471],
         [0.4969, 0.3340, 0.2211],
         [0.1553, 0.4028, 0.2000]]])

I found that the output of BatchNorm is not what I expected to be. For example, the mean across batch for first plane, first feature = 0.2518 and the std is 0.1572. The normalized value for the first value = (0.2961-0.2518)/0.1572 = 0.2818 != -0.5621.
My questions:

Am I correct to calculate the means in this way (across batch, per plane and feature)? as I understand batchnorm is used to treat the issue of having different scales for different feature, so it should at least be per feature dimension, however I am not sure whether to sum across the &quot;plane dimension&quot; as well.

Any other modifications I need to do to get the same output from BatchNorm2d?



This is the implementation of BatchNorm2d in pytorch (source1, source2). Using this, you can verify the operations you performed.
class MyBatchNorm2d(nn.BatchNorm2d):
    def __init__(self, num_features, eps=1e-5, momentum=0.1,
                 affine=True, track_running_stats=True):
        super(MyBatchNorm2d, self).__init__(
            num_features, eps, momentum, affine, track_running_stats)

    def forward(self, input):
        self._check_input_dim(input)

        exponential_average_factor = 0.0

        if self.training and self.track_running_stats:
            if self.num_batches_tracked is not None:
                self.num_batches_tracked += 1
                if self.momentum is None:  # use cumulative moving average
                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                else:  # use exponential moving average
                    exponential_average_factor = self.momentum

        # calculate running estimates
        if self.training:
            mean = input.mean([0, 2, 3])
            # use biased var in train
            var = input.var([0, 2, 3], unbiased=False)
            n = input.numel() / input.size(1)
            with torch.no_grad():
                self.running_mean = exponential_average_factor * mean\
                    + (1 - exponential_average_factor) * self.running_mean
                # update running_var with unbiased var
                self.running_var = exponential_average_factor * var * n / (n - 1)\
                    + (1 - exponential_average_factor) * self.running_var
        else:
            mean = self.running_mean
            var = self.running_var

        input = (input - mean[None, :, None, None]) / (torch.sqrt(var[None, :, None, None] + self.eps))
        if self.affine:
            input = input * self.weight[None, :, None, None] + self.bias[None, :, None, None]

        return input

The outputs of nn.BatchNorm2d(2)(a) and MyBatchNorm2d(2)(a) are same.
