28批量规范化

"""可持续加速深层网络的收敛速度"""
import torch
from torch import nn
import liliPytorch as lp
import matplotlib.pyplot as plt

def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
    """实现一个具有张量的批量规范化层。"""
    # 如果启用了梯度计算，torch.is_grad_enabled() 返回 True；否则返回 False。
    if not torch.is_grad_enabled():
        # torch.no_grad() 是一个上下文管理器，用于临时禁用梯度计算
        # torch.enable_grad() 是一个上下文管理器，用于在禁用梯度计算的上下文中重新启用梯度计算。
        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) # 计算张量 X 沿着第 0 维的平均值
            # 维度 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
    # gamma 和 beta 的更新是通过反向传播和优化器自动完成的
    Y = gamma * X_hat + beta # 缩放和移位
    return Y, moving_mean.data, moving_var.data

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)
        # 参与求梯度和迭代的拉伸和偏移参数，分别初始化成1和0
        self.gamma = nn.Parameter(torch.ones(shape))
        self.beta = nn.Parameter(torch.zeros(shape))
        
        # 非模型参数的变量初始化为0和1
        # 经过归一化处理后的数据均值接近于零。因此，将滑动均值初始化为0，是对数据初始均值的一种合理假设。
        self.moving_mean = torch.zeros(shape)
        # 方差表示数据的离散程度。将滑动方差初始化为1，意味着假设数据的初始方差为1，
        # 即数据分布接近标准正态分布。这样初始化可以避免初始阶段的数值不稳定。
        self.moving_var = torch.ones(shape)

    def forward(self, X):
        # 如果X不在内存上，将moving_mean和moving_var
        # 复制到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)
        
         # 保存更新过的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

#使用批量规范化层的 LeNet
net = nn.Sequential(
    nn.Conv2d(1, 6,  kernel_size=5, padding=2), # 卷积层1：输入通道数1，输出通道数6，卷积核大小5x5，填充2
    BatchNorm(num_features=6, num_dims=4),
    nn.ReLU(), # 激活函数
    nn.AvgPool2d(kernel_size=2, stride=2), # 平均池化层1：池化窗口大小2x2，步幅2

    nn.Conv2d(6, 16, kernel_size=5), # 卷积层2：输入通道数6，输出通道数16，卷积核大小5x5
    BatchNorm(num_features=16, num_dims=4),
    nn.ReLU(), 
    nn.AvgPool2d(kernel_size=2, stride=2), # 平均池化层2：池化窗口大小2x2，步幅2

    nn.Flatten(), # 展平层：将多维输入展平为1维
    nn.Linear(16 * 5 * 5, 120), # 全连接层1：输入节点数16*5*5，输出节点数120
    BatchNorm(num_features=120, num_dims=2),
    nn.ReLU(),
    nn.Linear(120, 84), # 全连接层2：输入节点数120，输出节点数84
    BatchNorm(num_features=84, num_dims=2),
    nn.ReLU(), 
    nn.Linear(84, 10) # 全连接层3：输入节点数84，输出节点数10（对应10个分类）
)

lr, num_epochs, batch_size = 1.0, 10, 256
train_iter, test_iter = lp.loda_data_fashion_mnist(batch_size)
# lp.train_ch6(net, train_iter, test_iter, num_epochs, lr, lp.try_gpu())
# plt.show()

# loss 0.200, train acc 0.925, test acc 0.812
# 34957.3 examples/sec on cuda:0

# loss 0.189, train acc 0.928, test acc 0.894
# 33471.2 examples/sec on cuda:0


#简明实现
net = nn.Sequential(
    nn.Conv2d(1, 6, kernel_size=5), nn.BatchNorm2d(6), nn.ReLU(),
    nn.AvgPool2d(kernel_size=2, stride=2),
    nn.Conv2d(6, 16, kernel_size=5), nn.BatchNorm2d(16), nn.ReLU(),
    nn.AvgPool2d(kernel_size=2, stride=2), nn.Flatten(),
    nn.Linear(256, 120), nn.BatchNorm1d(120), nn.ReLU(),
    nn.Linear(120, 84), nn.BatchNorm1d(84), nn.ReLU(),
    nn.Linear(84, 10)
)
lp.train_ch6(net, train_iter, test_iter, num_epochs, lr, lp.try_gpu())
plt.show()

# nn.Sigmoid()
# loss 0.263, train acc 0.902, test acc 0.833
# 46935.0 examples/sec on cuda:0

# nn.ReLU()
# loss 0.224, train acc 0.914, test acc 0.874
# 44479.2 examples/sec on cuda:0
"""
通常高级API变体运行速度快得多，因为它的代码已编译为C++或CUDA，而我们的自定义代码由Python实现。
"""