【李沐深度学习】59.BN层实线 代码示例

具体有关Batch Normlization的解释 可阅读blog

import torch
from torch import nn

'''
bn操作训练和推理时用的均值方差是不同的：
    - 训练时，是当前小批量中图片的均值方差，不同批次是不一样的；
    - 推理时，是全局样本的均值方差，固定的；
     
gamma, beta : bn层可学习的参数；
moving_mean, moving_var : 推理时用到的 全局均值和方差；
eps: 方差必须要加的极小值，防止归一化时除以的分母为0 导致出错；
momentum : 更新moving_mean,moving_var的参数，通常取0.9或其他固定值；
'''
def batch_norm(x, gamma, beta, moving_mean, moving_var, eps, momentum):
    if not torch.is_grad_enabled(): # 若是推理时(梯度记录未开启)
        x_hat = (x - moving_mean) / torch.sqrt(moving_var + eps)
    else: # 训练时
        assert len(x.shape) in (2, 4) # 输入特征图x的shape 为2，则是全连接层输出； 为4，则是2D卷积层输出；
        if len(x.shape) == 2:
            mean = x.mean(dim=0) # n*c dim=0 指 按行压缩，保留列---> 1*c； shape(5,3)--->(1,3)
            var = ((x - mean)**2).mean(dim=0)
        else:
            mean = x.mean(dim=(0,2,3), keepdim=True) # nchw，保留通道数，其他都算均值---> 因为keepdim所以还是4维 1*c*1*1
            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_mean + (1.0 - momentum) * var

    y = gamma * x_hat + beta
    return y, moving_mean.data, moving_var.data # 如果moving_mean/var是tensor的话，返回.data即可。梯度diff没用

# ------定义一个bn层------
class BatchNorm(nn.Module):
    def __init__(self, num_features, num_dims):
        super().__init__()
        if num_dims == 2: # 根据输入特征图形状，确定参数的形状
            shape = (1, num_dims)
        else:
            shape = (1, num_dims, 1, 1)
        self.gamma = nn.Parameter(torch.ones(shape)) # 注意是可学习参数，用nn.Parameter()包着
        self.beta = nn.Parameter(torch.zeros(shape))
        self.moving_mean = torch.zeros(shape)
        self.moving_var = torch.ones(shape)

    def forward(self, x):
        # 因为nn.Parameter()管理的可学习参数，在训练时xx.to(device)会自动转移到device对应的设备上，如GPU。此处moving_mean/var要自己管理
        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


# 若是直接掉包的话,网络示例如下：
net = nn.Sequential(
    nn.Conv2d(1, 6, kernel_size=5),
    nn.BatchNorm2d(6), # 上层输出的通道数
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),

    nn.Conv2d(6, 16, kernel_size=5),
    nn.BatchNorm2d(16),
    nn.LeakyReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),

    nn.Flatten(),
    nn.Linear(256, 120)
    nn.BatchNorm2d(120),

    nn.ReLU()
    nn.Linear(120, 10)
)