torch.nn.BatchNorm2d()计算

from torch.nn import MSELoss
import torch

class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.convs = torch.nn.ModuleList([
            torch.nn.BatchNorm2d(2, affine=True)
        ])

    def forward(self, x):
        for layer in self.convs:
            x = layer(x)
            print(x.shape)
        return x

data = torch.tensor([[[2, 1, 1], [1, 1, 1], [1, 1, 1]], [[1, 1, 1], [1, 1, 1], [1, 1, 1]]], dtype=torch.float32)
data = torch.unsqueeze(data, dim=0)
print(data.shape)
model = Model()
r = model(data)
print(r)
mean = data[0, 0].mean()
var = (data[0, 0] - data[0, 0].mean()).pow(2) / 8
norm = (data[0, 0] - mean) / (var.sqrt()+torch.tensor([1e-5]))
print(norm)
print("affine params", list(model.convs[0].named_parameters()))