深度学习笔记 —— 批量归一化

梯度在上面（损失处）的时候比较大，越到下面越容易变小，因为很多时候都是n个很小的数相乘，乘到最后梯度就比较小了。所以就导致上面参数更新快，而下面参数更新慢（下面参数在小范围内变化时，抽取的底层特征变化不大，此时上层的参数是针对这些底层特征进行学习的）。这也意味着，如果下面的参数改变了，那么上面的参数之前也就白学了，需要重新训练。使得收敛比较慢。

 核心想法：方差和均值的分布不同层之间变化，如果把分布固定住了，相对来说就是比较稳定的。模型稳定了，也就是说更新的时候不会爆炸，也不会太小，收敛就不会变慢。

批量归一化是个线性变换，把均值、方差拉动得比较好，使其变化不那么剧烈。

 此解释是否为正确的理论，也不一定……

允许用更大的学习率来做训练（使得梯度的值变大一点，每层之间梯度的值会差不多一点，所以可以用更大的学习率，对权重的更新变快）

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


def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
    # 不算梯度，即在做inference
    # 此处用的是全局的均值和方差，因为做inference的时候可能只有一个样本，算不出来批量的均值和方差
    if not torch.is_grad_enabled():
        X_hat = (X - moving_mean) / torch.sqrt(moving_var + eps)
    else:
        assert len(X.shape) in (2, 4)
        # 全连接，第一维是批量大小，第二维是特征
        if len(X.shape) == 2:
            mean = X.mean(dim=0)  # 对每一列算均值，mean是1xn的行向量
            var = ((X - mean) ** 2).mean(dim=0)  # 同理
        # 2D卷积，批量大小x通道数x高x宽
        else:
            mean = X.mean(dim=(0, 2, 3), keepdim=True)  # 对每一个通道，把所有批量、高、宽里面的像素求均值，结果是1xnx1x1的向量
            var = ((X - mean) ** 2).mean(dim=(0, 2, 3), keepdim=True)  # 1xnx1x1
        X_hat = (X - mean) / torch.sqrt(var + eps)
        # 更新moving_mean和moving_var
        moving_mean = momentum * moving_mean + (1.0 - momentum) * mean
        moving_var = momentum * moving_var + (1.0 - momentum) * var
    Y = gamma * X_hat + beta
    return Y, moving_mean.data, moving_var.data


# 创建一个正确的BatchNorm图层
class BatchNorm(nn.Module):
    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.zeros(shape))
        self.moving_mean = torch.zeros(shape)
        self.moving_var = torch.ones(shape)

    def forward(self, X):
        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)
        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


# 应用BatchNorm于LeNet模型
net = nn.Sequential(
    nn.Conv2d(1, 6, kernel_size=5), BatchNorm(6, num_dims=4), nn.Sigmoid(),
    nn.AvgPool2d(kernel_size=2, stride=2),
    nn.Conv2d(6, 16, kernel_size=5), BatchNorm(16, num_dims=4), nn.Sigmoid(),
    nn.AvgPool2d(kernel_size=2, stride=2), nn.Flatten(),
    nn.Linear(16*4*4, 120), BatchNorm(120, num_dims=2), nn.Sigmoid(),
    nn.Linear(120, 84), BatchNorm(84, num_dims=2), nn.Sigmoid(),
    nn.Linear(84, 10))


lr, num_epochs, batch_size = 1.0, 10, 256
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())

# 拉伸参数gamma和偏移参数beta
print(net[1].gamma.reshape((-1,)), net[1].beta.reshape((-1,)))


# 简明实现
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))