论文要点：BatchNorm（模型M+公式F+代码C）

# batchnorm

1. 代码

import torch
from torch import nn
from d2l import torch as d2l、

# moving_mean, moving_var 近似看作整个数据集上的均值和方差，
# eps是epsilon, momentum（通常=0.9）用来更新moving_mean, moving_var 
def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
    # 通过is_grad_enabled来判断当前模式是训练模式（enabled），还是预测模式（not is_grad_enabled）
    if not torch.is_grad_enabled():
        # 如果是在预测模式下，直接使用传入的moving_mean, moving_var做标准化
        X_hat = (X - moving_mean) / torch.sqrt(moving_var + eps)
    else:
        # X.shape=2 是全连接层，=4是2D卷积层
        assert len(X.shape) in (2, 4)
        if len(X.shape) == 2:
            # 若为全连接，按列求特征的均值和方差
            mean = X.mean(dim=0)
            var = ((X - mean) ** 2).mean(dim=0)
        else:
            # 若2D卷积层， 计算axis=1的均值和方差
            # X 保持形状以做broadcasting，（batch:0, channel:1, hight:2, width:3）
            mean = X.mean(dim(0, 2, 3), keepdim=True) # 对channel求均值
            var = ((X - mean) ** 2).mean(dim=(0, 2, 3), keepdim=True)
        # 训练模式下，用当前mean和var做标准化
        X_hat = (X - mean) / torch.sqrt(var + eps)
        # 不知共多少个，故用momentum求mean，var
        moving_mean = momentum * moving_mean + (1.0 - momentum) * mean
        moving_var = momentum * moving_var + (1.0 - momentum) * var
    # Y是X的线性变换，等同于X缩放和移位
    Y = gamma * X_hat +beta
    # 若torch.parameter, 则不要grad
    return Y, moving_mean.data, moving_var.data

# 创建BatchNorm层
class BatchNorm(nn.Module):
    # num_features：全连接层的输出数量或卷积层的输出通道数
    # num_dims：2表全连接，4表卷积
    def __init__(self, num_features, num_dims):
        super().__init__()
        if num_dims == 2:
            shape = (1, num_features)
        else:
            shape = (1, num_features, 1, 1)
        # gamma和beta是需要迭代的，所以放在nn.Parameter里
        self.gamma = nn.Parameter(torch.ones(shape))
        self.beta = nn.Parameter(torch.ones(shape))
        # 因未放在nn.Parameter，故需要自己算device
        self.moving_mean = torch.zeros(shape)
        self.moving_var = torch.ones(shape)
        
    def forward(self, X):
        # 复制到X所在GPU上
        if self.moving_mean.device != X.device:
            self.moving_mean = self.moving_mean.to(X.device)
            self.moving_var = self.moving_var.to(X.device)
        # 不同框架下eps值不同
        Y, self.moving_mean, self.moving_var = batch_norm(X, self.gamma, self.beta,
                                                         self.moving_mean, self.moving_var,
                                                         eps=1e-5, momentum=0.9)
        return Y

lr, num_epochs, batch_size = 1.0, 10, 1256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu)

net[1].gamma.reshape((-1,)), net[1].beta.reshape((-1,))

# 调用nn.BatchNorm2d()
net = nn.Sequential(
    nn.Conv2d(1, 6, kernel_size=5), nn.BatchNorm2d(6), nn.Sigmoid(),
    nn.AvgPool2d(kernel_size=2, stride=2),
    nn.Conv2d(6, 16, kernel_size=5), nn.BatchNorm2d(16), nn.Sigmoid(),
    nn.AvgPool2d(kernel_size=2, stride=2), nn.Flatten(),
    nn.Linear(256, 120), nn.BatchNorm1d(120), nn.Sigmoid(),
    nn.Linear(120, 84), nn.BatchNorm1d(84), nn.Sigmoid(),
    nn.Linear(84, 10))