CNN批量归一化实战：解决训练不稳定与梯度消失的终极方案

1次阅读

没有评论

共计 2174 个字符，预计需要花费 6 分钟才能阅读完成。

在深度神经网络中，第 $l$ 层的输入分布会随着前一层参数更新而改变，这种现象称为内部协变量偏移（Internal Covariate Shift）。数学表达为：

$$H_l = W_l X_{l-1} + b_l$$

其中 $X_{l-1}$ 是前层输出，其分布变化会导致当前层需要不断适应新的输入分布，进而引发：

梯度消失 / 爆炸：各层输入尺度差异迫使使用更小的学习率
饱和非线性：Sigmoid 等激活函数在极端值区域梯度接近于零

传统归一化方法（如输入特征标准化）仅作用于网络最底层，无法解决深度网络中的分层分布漂移问题。

批量归一化通过在每个小批量数据上执行标准化操作来解决上述问题，具体分为四个步骤：

计算批量统计量：
$$\mu_B = \frac{1}{m}\sum_{i=1}^m x_i$$
$$\sigma_B^2 = \frac{1}{m}\sum_{i=1}^m (x_i – \mu_B)^2$$
标准化处理：
$$\hat{x}_i = \frac{x_i – \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}$$
仿射变换：
$$y_i = \gamma \hat{x}_i + \beta$$

关键设计点：

$\gamma$/$\beta$ 作为可学习参数，保留网络的表示能力
$\epsilon$（默认 1e-5）防止除以零
推理阶段使用全局统计量替代批量统计量

import torch
import torch.nn as nn

class BatchNormFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, gamma, beta, running_mean, running_var, eps=1e-5, momentum=0.1):
        # 训练模式
        if x.requires_grad:
            dims = [d for d in range(x.dim()) if d != 1]
            mean = x.mean(dim=dims)
            var = x.var(dim=dims, unbiased=False)

            # EMA 更新全局统计量
            running_mean.mul_(1 - momentum).add_(mean * momentum)
            running_var.mul_(1 - momentum).add_(var * momentum)

            # 标准化
            x_hat = (x - mean.view(1, -1, 1, 1)) / torch.sqrt(var.view(1, -1, 1, 1) + eps)
            ctx.save_for_backward(x_hat, gamma, var)
        else:
            # 推理模式
            x_hat = (x - running_mean.view(1, -1, 1, 1)) / torch.sqrt(running_var.view(1, -1, 1, 1) + eps)

        return gamma.view(1, -1, 1, 1) * x_hat + beta.view(1, -1, 1, 1)

    @staticmethod
    def backward(ctx, grad_output):
        x_hat, gamma, var = ctx.saved_tensors

        # 根据论文公式 (6) 计算梯度
        grad_x = gamma / torch.sqrt(var + ctx.eps) * (
            grad_output 
            - x_hat * (grad_output * x_hat).mean(dim=(0,2,3), keepdim=True) 
            - grad_output.mean(dim=(0,2,3), keepdim=True)
        )

        grad_gamma = (grad_output * x_hat).sum(dim=(0,2,3))
        grad_beta = grad_output.sum(dim=(0,2,3))

        return grad_x, grad_gamma, grad_beta, None, None, None, None

测试环境：RTX 3090, CUDA 11.3, PyTorch 1.10