深入理解BatchNorm反向传播：从数学推导到PyTorch实现避坑指南

1次阅读

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

BatchNorm 在训练和推理阶段的行为差异是第一个需要理解的要点。训练时，它用当前 batch 的均值 / 方差做归一化：

$$\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$$

而在 eval 模式时，则使用全局统计量 running_mean/running_var。这种差异导致手动实现时容易忽略对 $\mu_B$ 和 $\sigma_B^2$ 的梯度计算——它们是前向传播时的中间变量，但反向传播时也需要参与链式法则。

计算 batch 均值 $\mu_B$
计算 batch 方差 $\sigma_B^2$
归一化：$\hat{x}_i = \frac{x_i-\mu_B}{\sqrt{\sigma_B^2+\epsilon}}$
缩放平移：$y_i = \gamma\hat{x}_i + \beta$

对输入 $x_i$ 的梯度包含三部分：

$$\frac{\partial L}{\partial x_i} = \frac{\partial L}{\partial y_i} \cdot \gamma \cdot \left(\frac{1}{\sqrt{\sigma_B^2+\epsilon}} – \frac{(x_i-\mu_B)^2}{m(\sigma_B^2+\epsilon)^{3/2}}\right) + \frac{\partial L}{\partial \mu_B}\cdot\frac{-1}{m} + \frac{\partial L}{\partial \sigma_B^2}\cdot\frac{-2(x_i-\mu_B)}{m}$$

其中对 $\mu_B$ 和 $\sigma_B^2$ 的梯度常被忽略，这是手动实现的主要难点。完整推导建议参考原论文的补充材料。

import torch
import torch.nn as nn

class ManualBatchNorm2d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(num_features, 1, 1))
        self.beta = nn.Parameter(torch.zeros(num_features, 1, 1))
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
        self.eps = eps
        self.momentum = momentum

    def forward(self, x):
        if self.training:
            # 计算 batch 统计量
            mean = x.mean(dim=(0, 2, 3), keepdim=True)
            var = x.var(dim=(0, 2, 3), unbiased=False, keepdim=True)

            # 更新 running 统计量
            self.running_mean = (1-self.momentum)*self.running_mean + self.momentum*mean.squeeze()
            self.running_var = (1-self.momentum)*self.running_var + self.momentum*var.squeeze()

            # 归一化
            x_hat = (x - mean) / torch.sqrt(var + self.eps)
        else:
            x_hat = (x - self.running_mean.view(1,-1,1,1)) / \
                   torch.sqrt(self.running_var.view(1,-1,1,1) + self.eps)

        return self.gamma * x_hat + self.beta

class AutoBatchNorm2d(nn.Module):
    def __init__(self, num_features):
        super().__init__()
        self.bn = nn.BatchNorm2d(num_features, affine=False)
        self.gamma = nn.Parameter(torch.ones(num_features, 1, 1))
        self.beta = nn.Parameter(torch.zeros(num_features, 1, 1))

    def forward(self, x):
        return self.gamma * self.bn(x) + self.beta