47. 批量规范化 代码从零开始实现 + 简洁实现

1. 从零实现

下面，我们从头开始实现一个具有张量的批量规范化层。

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

# X是输入，gamma, beta是两个可以学习的参数，moving_mean, moving_var是整个数据集的均值和方差，
# 而不是小批量的均值和方差，做inference（test）的时候用，eps防止除零错误，通常不改变值
# momentum是来更新moving_mean, moving_var的
def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
    # 通过is_grad_enabled来判断当前模式是训练模式还是预测模式
    # 训练和预测时 BN的过程不一样
    if not torch.is_grad_enabled(): # 不用计算梯度，就代表不是在做训练，是做inference
        # 推理就是forward，不需要反向传播，所以不需要计算梯度
        # 如果是在预测模式下，直接使用传入的移动平均所得的均值和方差
        X_hat = (X - moving_mean) / torch.sqrt(moving_var + eps) # 推理时用的全局的mean和var
        # 为什么不能用一个批量的均值和方差？
        # 答：因为做inference的时候，可能没有一个批量，可能就一个样本进来

    else: # 训练模式下
        assert len(X.shape) in (2, 4) # X.shape要么为2，表示全连接层，要么为4，表示卷积层
        if len(X.shape) == 2:# (batch_size,feature)
            # 使用全连接层的情况，计算特征维上的均值和方差
            mean = X.mean(dim=0) # 表示跨行（dim=0）求平均，也就是对同一列元素求均值
            var = ((X - mean) ** 2).mean(dim=0) # 因此方差和mean都是一个行向量

        else: # （batch_size,channels,height,width）
            # 使用二维卷积层的情况，计算通道维上（axis=1）的均值和方差。
            # 这里我们需要保持X的形状以便后面可以做广播运算
            mean = X.mean(dim=(0, 2, 3), keepdim=True) 
            # dim=(0, 2, 3)表示对每一个通道的所有批量*所有高宽的像素全部来求均值
            # 换个角度想：dim=1的话表示哭啊通道求平均，也就是每一行求平均
            # dim=(0, 2, 3)与之相反，那么就是每一列求平均
            # 最后求出来的结果是 1 * n * 1 * 1，因为keepdim=True
            var = ((X - mean) ** 2).mean(dim=(0, 2, 3), keepdim=True)
            # X - mean用了广播机制，X为2 *2 * 3* 3，mean为 1 * 2 * 1 * 1，可以进行广播机制
            # 最后求出来的结果也是 1 * n * 1 * 1
        # 训练模式下，用当前的均值和方差做标准化
        X_hat = (X - mean) / torch.sqrt(var + eps) # 训练时使用的是当前的小批量中的mean和方差

        # 更新移动平均的均值和方差：
        # 每得到一次小批量的均值和方差，就去更新全局的数据，这样会无限逼近真实的在数据集上的均值和方差
        moving_mean = momentum * moving_mean + (1.0 - momentum) * mean
        moving_var = momentum * moving_var + (1.0 - momentum) * var

    # 返回的Y就是对X做了线性变化之后的数据
    Y = gamma * X_hat + beta  # 缩放和移位
    return Y, moving_mean.data, moving_var.data 

我们现在可以创建一个正确的BatchNorm层。 这个层将保持适当的参数：拉伸gamma和偏移beta,这两个参数将在训练过程中更新。 此外，我们的层将保存均值和方差的移动平均值，以便在模型预测期间随后使用。

撇开算法细节，注意我们实现层的基础设计模式。

通常情况下，我们用一个单独的函数定义其数学原理，比如说batch_norm。 然后，我们将此功能集成到一个自定义层中，其代码主要处理数据移动到训练设备（如GPU）、分配和初始化任何必需的变量、跟踪移动平均线（此处为均值和方差）等问题。

为了方便起见，我们并不担心在这里自动推断输入形状，因此我们需要指定整个特征的数量。 不用担心，深度学习框架中的批量规范化API将为我们解决上述问题，我们稍后将展示这一点。

class BatchNorm(nn.Module):
    # num_features：完全连接层的输出数量或卷积层的输出通道数。
    # num_dims：2表示完全连接层，4表示卷积层
    def __init__(self, num_features, num_dims):
        super().__init__()
        # 为什么要定义shape的原因？
        # 答：gamma，beta，moving_mean和moving_var的shape都是和shape一样
        if num_dims == 2: # 完全连接层
            shape = (1, num_features)
        else: # 卷积层
            shape = (1, num_features, 1, 1)
        # 参与求梯度和迭代的拉伸和偏移参数，分别初始化成1和0
        # gamma如果初始化为0，那么gamma乘以X_hat等于0，训练不动
        self.gamma = nn.Parameter(torch.ones(shape)) 
        self.beta = nn.Parameter(torch.zeros(shape))
        # gamma和 beta是需要被迭代的，moving_mean和moving_var则不需要被迭代
        # 所以gamma和 beta放在nn.Parameter中，moving_mean和moving_var就不需要放在里面
        # 放入nn.Parameter中的变量，会随着网络的优化器一起更新迭代

        # 非模型参数的变量初始化为0和1
        self.moving_mean = torch.zeros(shape)
        self.moving_var = torch.ones(shape)

    def forward(self, X):
        # 如果X不在内存上，将moving_mean和moving_var
        # 复制到X所在显存上
        if self.moving_mean.device != X.device:
          # 因为moving_mean和moving_var 没有放在nn.Parameter中，也就是没放入网络net中，
          # 而网络是在GPU上，所以需要移动这两个参数
            self.moving_mean = self.moving_mean.to(X.device)
            self.moving_var = self.moving_var.to(X.device)
        # 保存更新过的moving_mean和moving_var
        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

2. 使用批量规范化层的 LeNet

为了更好理解如何应用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))

和以前一样，我们将在Fashion-MNIST数据集上训练网络。 这个代码与我们第一次训练LeNet时几乎完全相同，主要区别在于学习率大得多。

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。

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

3. 简洁实现

除了使用我们刚刚定义的BatchNorm，我们也可以直接使用深度学习框架中定义的BatchNorm。 该代码看起来几乎与我们上面的代码相同。

net = nn.Sequential(
    # nn.BatchNorm2d(6)传入的参数是：完全连接层的输出数量或卷积层的输出通道数
    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))

下面，我们使用相同超参数来训练模型。 请注意，通常高级API变体运行速度快得多，因为它的代码已编译为C++或CUDA，而我们的自定义代码由Python实现。

d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())

4. 一些Q&A

Q1: 之前讲xavier的时候，也讲过类似的normalization的东西，这个和这里的BN有什么区别？

A1:讲xavier是选取好的初始化，对权重进行好的初始化（对权重分布的均值和方差进行控制），使模型训练比较稳定，但不能保证之后。

而BN则是保证在整个模型训练的时候对feature map进行均值和方差的控制，即：在每一个全连接层和卷积层前面或后面做normalization。

Q2：批量归一化是不是和权重衰退的作用类似？

A2：不类似。权重衰退是指在更新参数（权重）的时候，把权重除了一个值，使其变小，但批量归一化是不会对前面一层的权重有影响，他只是对上一层输出的数据进行规范化（均值和方差的控制）。

Q3: batch norm 能用在mlp中吗？

A3:可以的，但是batch norm对比较深的神经网络相对有用，因为只有当网络比较深的时候，才会出现上面更新快，而靠近输入数据的更新慢的情况，此时用BN效果会好。浅层MLP加上BN效果不见得很好。

Q4: BN是做了线性变换，和加一个线性层有什么区别？

A4：如果我就做一般的线性变换，那我对这个数据做怎样的处理呢？我也不知道要这个线性层做什么才能使得数值稳定，我加一个线性层能学到的参数不一定能学得这么好。而BN就相当于一个先验知识，告诉你这么做是可以使得训练更稳定的。并且，一般的线性变化Linear是同一个样本的各个维度，而BN是针对不同样本的同一特征。

Q5：为什么加了BN之后，收敛时间变短？

A5：BN会使得梯度的值变大一些，而且每个层的梯度的值会差不多，即不同层之间的梯度稳定，这样导致可以使用更大的学习率，就会使得对权重的更新变快。

有一种解释是，BN之后地形图更加平坦，更容易收敛。

Q6：layer normalization和 batch normalization 的异同。

A6： batch normalization是在feature维度对样本做normalization（对不同样本的同一特征维度做规范化）。

layer normalization就不是像batch normalization跨样本做规范化，而是在每一个样本里面做规范化，也就是对同一个样本的所有特征进行规范化。

Q7：batch size 是把显存占满好？还是利用率gpu-util 100%就可以了？还是需要同时达到100%？

A7:增加bath size ，去看一下每一秒能处理的样本数，如果增到一个程度发现每一秒能处理的样本数不会再增加，就可以停止增加batch size。但这不针对LeNet，因为LeNet可以有很大的batch size，增大到1000都可以，但是收敛会有问题。