Batch Normalization实战指南：如何解决深度神经网络训练中的梯度消失问题

1次阅读

没有评论

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

在深度神经网络训练过程中，梯度消失（Gradient Vanishing）和内部协变量偏移（Internal Covariate Shift）是两个常见的痛点。具体表现为：

梯度消失：随着网络层数加深，反向传播时梯度会逐层衰减，导致浅层参数几乎无法更新。
内部协变量偏移：每层输入的分布会随着前一层参数更新而不断变化，迫使后续层必须频繁适应新的数据分布。

传统解决方案如使用 ReLU 激活函数、精心初始化权重（Xavier/Glorot 初始化）等，只能部分缓解问题。而 Batch Normalization（BN）通过标准化每一层的输入分布，从根本上改善了这两个问题。

给定一个 mini-batch 输入 $B = {x_1, …, x_m}$，BN 层执行以下操作：

计算 mini-batch 均值：
$$\mu_B = \frac{1}{m}\sum_{i=1}^m x_i$$
计算 mini-batch 方差：
$$\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$$

γ（scale）：允许网络决定标准化后的缩放程度
β（shift）：允许网络决定标准化后的偏移量
这两个参数让网络可以学习是否使用 BN 带来的标准化效果。

import torch
import torch.nn as nn

class CustomBatchNorm1d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        self.eps = eps
        self.momentum = momentum

        # 可训练参数
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))

        # 非训练参数（推理时使用）self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))

    def forward(self, x):
        if self.training:
            # 训练模式：使用当前 batch 统计量
            mean = x.mean(dim=0)
            var = x.var(dim=0, unbiased=False)

            # 更新 running 统计量
            with torch.no_grad():
                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
        else:
            # 推理模式：使用 running 统计量
            mean = self.running_mean
            var = self.running_var

        # 标准化
        x_hat = (x - mean) / torch.sqrt(var + self.eps)

        # 缩放和平移
        return self.gamma * x_hat + self.beta

class CNNWithBN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3)
        self.bn1 = nn.BatchNorm2d(16)  # 注意 Conv2d 对应 BatchNorm2d
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3)
        self.bn2 = nn.BatchNorm2d(32)
        self.fc = nn.Linear(32*6*6, 10)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        return self.fc(x)

我们在 CIFAR-10 数据集上对比了使用 BN 和不使用 BN 的 ResNet-18 模型：