import torch 
import torch.nn as nn 
batch, ch, h, w = 2, 32, 128, 128
torch_batch_norm = nn.BatchNorm2d(num_features=ch)
for k, v in torch_batch_norm.named_parameters():
    print(f'parameter: {k}', v.shape)
for k, v in torch_batch_norm.named_buffers():
    print(f'buffer: {k}', v.shape)

import torch
import torch.nn as nn 
class MyBatchNorm(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        # weight初始化为1，bias初始化为0.
        self.weight = nn.Parameter(torch.ones(num_features))
        self.bias = nn.Parameter(torch.zeros(num_features))

        self.momentum = 0.1 
        self.eps = eps 

        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
        self.register_buffer('num_batches_tracked', torch.tensor(0.))

    def forward(self, tensor): 
        bs, ch, h, w = tensor.shape 
        # 变换一下tensor的尺寸，方便处理
        tensor_flatten = tensor.permute(0, 2, 3, 1).flatten(0, 2)  # bs * h * w, ch 
        # 求均值和方差
        mean = torch.mean(tensor_flatten, 0)
        # 注意方差有biased和unbiased两种。
        var = torch.var(tensor_flatten, 0, unbiased=False)
        var_unbiased = torch.var(tensor_flatten, 0, unbiased=True)

        if self.training:
            # 训练时，我们要执行moving average，统计
            # running_mean和running_var，注意此时应
            # 使用unbiased版本的方差。
            self.running_mean.mul_(1 - self.momentum).add_(self.momentum * mean)
            self.running_var.mul_(1 - self.momentum).add_(self.momentum * var_unbiased)
            self.num_batches_tracked.add_(1)

        # 训练时用batch内的统计量，测试时用moving average
        # 保存的统计量。
        if self.training:
            tensor_flatten = (tensor_flatten - mean) / torch.sqrt(var + self.eps)
        else: 
            tensor_flatten = (tensor_flatten - self.running_mean) / torch.sqrt(self.running_var + self.eps)

        # 归一化完成后，做线性变换。
        ret = tensor_flatten * self.weight + self.bias 
        ret = ret.view(bs, h, w, ch).permute(0, 3, 1, 2)
        return ret 

my_batch_norm = MyBatchNorm(num_features=ch)
# 因为BN涉及running_mean和running_var的更新，所以我们要多跑几轮来检查moving average的正确性。
for _ in range(10):
    a = torch.rand(batch, ch, h, w)
    ret1 = torch_batch_norm(a)
    ret2 = my_batch_norm(a)
    diff = torch.mean(torch.abs(ret1 - ret2)).item()
    
    running_mean1 = torch_batch_norm.running_mean 
    running_mean2 = my_batch_norm.running_mean 
    diff_mean = torch.mean(torch.abs(running_mean1 - running_mean2)).item()

    running_var1 = torch_batch_norm.running_var 
    running_var2 = my_batch_norm.running_var
    diff_var = torch.mean(torch.abs(running_var1 - running_var2)).item()
    print('{:.6f};{:.6f};{:.6f}'.format(diff, diff_mean, diff_var))

# .eval()开启测试模型
torch_batch_norm.eval()  
my_batch_norm.eval()
for _ in range(10):
    a = torch.rand(batch, ch, h, w)
    ret1 = torch_batch_norm(a)
    ret2 = my_batch_norm(a)
    diff = torch.mean(torch.abs(ret1 - ret2)).item()
    
    running_mean1 = torch_batch_norm.running_mean 
    running_mean2 = my_batch_norm.running_mean 
    diff_mean = torch.mean(torch.abs(running_mean1 - running_mean2)).item()

    running_var1 = torch_batch_norm.running_var 
    running_var2 = my_batch_norm.running_var
    diff_var = torch.mean(torch.abs(running_var1 - running_var2)).item()
    print('{:.6f};{:.6f};{:.6f}'.format(diff, diff_mean, diff_var))