BN对抗梯度消失：从原理到PyTorch实战指南

1次阅读

没有评论

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

在训练深度神经网络时，梯度消失问题就像是一个隐形的障碍，阻碍着模型的有效学习。简单来说，梯度消失指的是在反向传播过程中，梯度信号随着网络层数的增加而指数级减小，导致浅层网络的权重几乎得不到更新。

这种现象在不同激活函数下表现各异：

Sigmoid 函数 ：其导数最大值为 0.25，经过多层连乘后梯度会迅速衰减
ReLU 函数 ：虽然正区间导数为 1，但负半轴的 ” 死亡 ReLU” 问题同样会导致部分神经元永久失效

面对梯度消失，开发者们尝试了多种解决方案，各有优劣：

权重初始化 （如 Xavier、He 初始化）
优点：实现简单，计算开销小
局限：只能缓解初始化阶段的梯度问题
层归一化（LayerNorm）
优点：对 batch 大小不敏感，适合 RNN 结构
局限：在 CNN 中效果不如 BN 显著
批归一化（BN）
优势：
- 允许使用更高学习率
- 减少对初始化的依赖
- 自带轻微正则化效果
缺点：
- 小 batch 下统计量不准确
- 增加计算和内存开销

批归一化的实现可以分为四个关键步骤：

计算 batch 均值：
$$\mu = \frac{1}{m}\sum_{i=1}^m x_i$$
计算 batch 方差：
$$\sigma^2 = \frac{1}{m}\sum_{i=1}^m (x_i-\mu)^2$$
归一化处理（加上 ε 防止除零）：
$$\hat{x} = \frac{x-\mu}{\sqrt{\sigma^2 + \epsilon}}$$
缩放平移（引入可学习参数 γ 和 β）：
$$y = \gamma \hat{x} + \beta$$

import torch
import torch.nn as nn

class BN_CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(64, eps=1e-5, momentum=0.1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)
        self.fc = nn.Linear(128*8*8, 10)

    def forward(self, x):
        x = torch.relu(self.bn1(self.conv1(x)))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.bn2(self.conv2(x)))
        x = torch.max_pool2d(x, 2)
        x = x.view(x.size(0), -1)
        return self.fc(x)

关键参数说明 ：

eps：防止方差为零的小常数（默认 1e-5）
momentum：更新 running 统计量的衰减系数（默认 0.1）

模式切换示例 ：

model = BN_CNN()

# 训练模式（默认）model.train()  
# 会更新 running_mean 和 running_var

# 推理模式
model.eval()
# 使用训练阶段积累的 running 统计量

小 batch size 问题 ：
现象：batch<16 时统计量波动大
解决方案：
- 使用 Group Normalization 替代
- 增大 batch size 或使用累积 BN
与 Dropout 的配合 ：
BN 本身有正则效果，可适当减少 Dropout 率
建议使用顺序：Conv → BN → ReLU → Dropout
模型导出注意事项 ：
确保推理时使用 model.eval()

导出 ONNX 时注意包含 running_mean/var

torch.onnx.export(model, dummy_input, "model.onnx", 
                input_names=["input"], 
                output_names=["output"],
                dynamic_axes={"input": {0: "batch"}, 
                             "output": {0: "batch"}})

在 CIFAR-10 数据集上的训练曲线显示：