batch-normalization的维度是按照mini-batch的维度进行的。也就是对于维度为[n, c ,h , w]的input来说,其对应的BN的λ参数的维度是 [1, c, 1, 1]
import torch
import torch.nn as nn
input = torch.rand(4, 3, 5, 5, dtype=torch.float32)
print(input[0, 0, 0, :]) # 输入
# 自己写的batch-normalization 减均值,除以方差
meanval = torch.mean(input, dim=[0, 2, 3], keepdim=True)
stdval = torch.std(input, dim=[0, 2, 3], keepdim=True, unbiased=False)
res_manual = input - meanval
res_manual = torch.div(res_manual, stdval)
print(res_manual[0, 0, :, :])
# 自带的BN
bn = nn.BatchNorm2d(input.shape[1], affine=False, momentum=0.0, eps=0.0, track_running_stats=False)
res = bn(input)
print(res[0, 0, :, :])