Batch Norm反向传播的实现原理与工程优化指南

1次阅读

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

Batch Normalization（批归一化）是深度学习中一项关键技术，通过对每层输入进行归一化处理，能够显著加速模型收敛并提升训练稳定性。然而，在实际应用中，Batch Norm 的反向传播实现往往被忽视，导致训练过程中出现梯度不稳定、训练震荡等问题。这些问题主要源于：

计算均值、方差时的数值稳定性问题
反向传播时梯度计算的复杂性
小批量数据下统计量估计不准确

Batch Norm 的前向传播过程可以分为以下几步：

计算 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$$

反向传播时需要计算对各个参数的梯度。根据链式法则，我们需要计算：

对缩放参数 γ 的梯度：
$$\frac{\partial L}{\partial \gamma} = \sum_{i=1}^m \frac{\partial L}{\partial y_i} \hat{x}_i$$
对平移参数 β 的梯度：
$$\frac{\partial L}{\partial \beta} = \sum_{i=1}^m \frac{\partial L}{\partial y_i}$$
对输入 x 的梯度计算最为复杂，需要经过完整的反向传播链：
$$\frac{\partial L}{\partial x_i} = \frac{\partial L}{\partial \hat{x}_i} \cdot \frac{1}{\sqrt{\sigma_B^2 + \epsilon}} + \frac{\partial L}{\partial \sigma_B^2} \cdot \frac{2(x_i – \mu_B)}{m} + \frac{\partial L}{\partial \mu_B} \cdot \frac{1}{m}$$

import torch
import torch.nn as nn

class CustomBatchNorm1d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super(CustomBatchNorm1d, self).__init__()
        self.num_features = num_features
        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)

            # 更新运行统计量
            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:
            # 推理阶段使用保存的运行统计量
            mean = self.running_mean
            var = self.running_var

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

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

        return out

import tensorflow as tf

class CustomBatchNorm(tf.keras.layers.Layer):
    def __init__(self, epsilon=1e-5, momentum=0.1, **kwargs):
        super(CustomBatchNorm, self).__init__(**kwargs)
        self.epsilon = epsilon
        self.momentum = momentum

    def build(self, input_shape):
        dim = input_shape[-1]

        # 可训练参数
        self.gamma = self.add_weight(name='gamma', shape=[dim],
            initializer='ones', trainable=True)
        self.beta = self.add_weight(name='beta', shape=[dim],
            initializer='zeros', trainable=True)

        # 运行时统计量
        self.moving_mean = self.add_weight(name='moving_mean', shape=[dim],
            initializer='zeros', trainable=False)
        self.moving_variance = self.add_weight(name='moving_variance', shape=[dim],
            initializer='ones', trainable=False)

    def call(self, inputs, training=None):
        if training:
            # 训练阶段使用当前 batch 的统计量
            mean = tf.reduce_mean(inputs, axis=0)
            variance = tf.math.reduce_variance(inputs, axis=0)

            # 更新运行统计量
            self.moving_mean.assign(self.momentum * mean + (1 - self.momentum) * self.moving_mean)
            self.moving_variance.assign(self.momentum * variance + (1 - self.momentum) * self.moving_variance)
        else:
            # 推理阶段使用保存的运行统计量
            mean = self.moving_mean
            variance = self.moving_variance

        # 归一化
        x_hat = (inputs - mean) / tf.sqrt(variance + self.epsilon)

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