（动手学习深度学习）第7章 批量规范化（Batch Normalization）

BN

总结

• 批量归一化固定小批量中的均值和方差，然后学习出适合的偏移和缩放。
• 可以加速收敛速度，但一般不改变模型精度。

BN代码手动实现

1. 导入相关库

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

2. 定义BN层

def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
    """
    :param X: 输入数据
    :param gamma: γ
    :param beta: β
    :param moving_mean: 全局均值
    :param moving_var: 全局方差
    :param eps: ε
    :param momentum: 冲量：用来更新或固定常量
    :return:输出数据, 全局均值, 全局方差
    """
    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)
            var = ((X-mean)**2).mean(dim=0)
        else:
            # 使用二维卷积层情况：计算通道维上（axis=1）的均值和方差
            # 这里需要保持X的形状以便后面可以做广播运算
            mean = X.mean(dim = (0, 2, 3), keepdim=True)
            var = ((X - mean)**2).mean(dim=(0, 2, 3), keepdim=True)
        # 训练模式下，用当前的均值和方差做标准化
        X_hat = (X - mean) / torch.sqrt(var + eps)
        # 更新移动平均的均值和方差
        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

class BatchNorm(nn.Module):

    def __init__(self, num_features, num_dims):
        """
        :param num_features: 全连接层的输出数量或卷积层的输出通道数
        :param num_dims: 2表示全连接层，4表示卷积层
        """
        super().__init__()
        if num_dims == 2:
            shape = (1, num_features)
        else:
            shape = (1, num_features, 1, 1)
        # 参与求梯度和迭代的拉伸参数和偏移参数，其分别初始化为1和0
        self.gamma = nn.Parameter(torch.ones(shape))
        self.beta = nn.Parameter(torch.zeros(shape))
        # 非模型参数的变量初始化为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:
            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

3. 应用BN与LeNet模型

net = nn.Sequential(
    nn.Conv2d(1, 6, kernel_size=5, padding=2), BatchNorm(6, num_dims=4), nn.Sigmoid(), # [1, 6, 28, 28]
    nn.AvgPool2d(2, stride=2),  # [1, 6, 14, 14]

    nn.Conv2d(6, 16, kernel_size=5), BatchNorm(16, num_dims=4), nn.Sigmoid(),  # [1, 16, 10, 10]
    nn.AvgPool2d(2, stride=2),  # [1, 16, 7, 7]

    nn.Flatten(),  # [1, 16*5*5]

    nn.Linear(16*5*5, 120), BatchNorm(120, num_dims=2), nn.Sigmoid(),  # [1, 400] -->[1, 120]
    nn.Linear(120, 84), BatchNorm(84, num_dims=2), nn.Sigmoid(),  # [1, 120] --> [1, 84]
    nn.Linear(84, 10)  # [1, 82] --> [1, 10]
)

4. 查看模型

X = torch.randn((1, 1, 28, 28))
for layer in net:
    X = layer(X)
    print(layer.__class__.__name__, 'output shape: \t', X.shape)

5. 加载数据集

batch_size = 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)

6. 训练模型

import time
lr, num_epochs = 1.0, 10
start = time.perf_counter()
d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())
end = time.perf_counter()
print("运行耗时 %.4f s" % (end-start))

7. 查看拉伸参数gamma和偏移参数beta

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

BN代码简洁实现

• 修改模型

net = nn.Sequential(
    nn.Conv2d(1, 6, kernel_size=5, padding=2),
    nn.BatchNorm2d(6), nn.Sigmoid(),  # [1, 6, 28, 28]
    nn.AvgPool2d(2, stride=2),  # [1, 6, 14, 14]

    nn.Conv2d(6, 16, kernel_size=5),
    nn.BatchNorm2d(16), nn.Sigmoid(),  # [1, 16, 10, 10]
    nn.AvgPool2d(2, stride=2),  # [1, 16, 5, 5]

    nn.Flatten(),  # [1, 16*5*5]

    nn.Linear(16 * 5 * 5, 120),
    nn.BatchNorm1d(120), nn.Sigmoid(),  # [1, 120]
    nn.Linear(120, 84),
    nn.BatchNorm1d(84), nn.Sigmoid(),  # [1, 84]
    nn.Linear(84, 10))  # [1, 10]

• 训练模型

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