深入解析BN批量归一化：原理、实现与训练加速实践

1次阅读

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

在深度神经网络训练中，随着网络层数的加深，每层输入的分布会逐渐发生变化，这种现象被称为 Internal Covariate Shift（ICS）。Sergey Ioffe 在其论文中通过公式描述了这一现象：

$$
\Delta w \propto \frac{\partial L}{\partial w} = \frac{\partial L}{\partial f} \cdot \frac{\partial f}{\partial w}
$$

其中，$f$ 是层的输出，$L$ 是损失函数。由于每层的输入分布不断变化，导致梯度更新方向不稳定，进而使得模型收敛缓慢且训练不稳定。

BN（Batch Normalization）：适用于 CV 任务，依赖于 batch 维度统计量，对 batch size 敏感。
LayerNorm：常用于 NLP 任务，对单个样本的通道维度进行归一化，不依赖 batch size。
InstanceNorm：主要用于风格迁移等任务，对每个样本的每个通道单独归一化。

γ/ β 参数学习 ：BN 层引入了两个可学习参数 γ（缩放）和 β（偏移），用于恢复网络的表达能力。公式如下：

$$
y = \gamma \cdot \hat{x} + \beta
$$

推理阶段 EMA：在训练阶段，BN 层会计算移动平均值（EMA）保存均值和方差，推理阶段直接使用预存值，避免依赖 batch 统计量。

import torch.nn as nn

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(out_channels)  # 初始化 BN 层

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)  # 训练时计算 batch 统计量，推理时使用预存 EMA
        return x

import tensorflow as tf

class ConvBlock(tf.keras.Model):
    def __init__(self, filters):
        super().__init__()
        self.conv = tf.keras.layers.Conv2D(filters, 3, padding='same')
        self.bn = tf.keras.layers.BatchNormalization()  # 自动处理训练 / 推理模式

    def call(self, inputs, training=False):
        x = self.conv(inputs)
        x = self.bn(x, training=training)  # training 参数控制统计量计算
        return x