深入理解BN层的反向传播：从数学推导到PyTorch实现

1次阅读

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

批量归一化（Batch Normalization，BN）是现代深度神经网络中的基础组件，但它的反向传播过程常被当作黑箱使用。本文将拆解 BN 层的数学本质，并通过 PyTorch 实现揭示其训练细节。

BN 层在训练时需动态计算批次统计量（均值 / 方差），这些中间变量会参与梯度计算
标准化操作（减均值除方差）使得梯度流需要复合函数求导法则
推理阶段的固定统计量与训练时的动态统计量存在模式差异

设输入张量 $X \in \mathbb{R}^{N\times C\times H\times W}$，对通道 $c$ 的计算步骤如下：

前向传播：
$$
\mu_c = \frac{1}{NHW}\sum_{n,h,w}X_{n,c,h,w}
$$
$$
\sigma_c^2 = \frac{1}{NHW}\sum_{n,h,w}(X_{n,c,h,w}-\mu_c)^2 + \epsilon
$$
$$
\hat{X}{n,c,h,w} = \frac{X
$$
$$
Y_{n,c,h,w} = \gamma_c \hat{X}_{n,c,h,w} + \beta_c
$$}-\mu_c}{\sqrt{\sigma_c^2}

反向传播（令 $\frac{\partial L}{\partial Y}$ 为上游梯度）：

对缩放参数 $\gamma$ 的梯度：
$$
\frac{\partial L}{\partial \gamma_c} = \sum_{n,h,w}\frac{\partial L}{\partial Y_{n,c,h,w}}\hat{X}_{n,c,h,w}
$$
对平移参数 $\beta$ 的梯度：
$$
\frac{\partial L}{\partial \beta_c} = \sum_{n,h,w}\frac{\partial L}{\partial Y_{n,c,h,w}}
$$
对输入 $X$ 的梯度（经过链式法则展开后）：
$$
\frac{\partial L}{\partial X_{n,c,h,w}} = \frac{\gamma_c}{\sqrt{\sigma_c^2}}\left[\frac{\partial L}{\partial Y_{n,c,h,w}} – \frac{1}{NHW}\left(\sum_{k,h,w}\frac{\partial L}{\partial Y_{k,c,h,w}} + \hat{X}{n,c,h,w}\sum\right)\right]
$$}\frac{\partial L}{\partial Y_{k,c,h,w}}\hat{X}_{k,c,h,w

class CustomBatchNorm2d(nn.Module):
    def __init__(self, num_features, eps=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))
        self.eps = eps
        # 缓存推理时使用的统计量
        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:
            # 训练模式使用当前批次统计量
            mean = x.mean(dim=[0,2,3], keepdim=True)
            var = x.var(dim=[0,2,3], unbiased=False, keepdim=True)

            # 更新 running 统计量（需停止梯度追踪）with torch.no_grad():
                self.running_mean = 0.9 * self.running_mean + 0.1 * mean.squeeze()
                self.running_var = 0.9 * self.running_var + 0.1 * var.squeeze()
        else:
            # 推理模式使用缓存统计量
            mean = self.running_mean.view(1,-1,1,1)
            var = self.running_var.view(1,-1,1,1)

        # 标准化计算
        x_hat = (x - mean.detach()) / torch.sqrt(var.detach() + self.eps)
        return self.gamma.view(1,-1,1,1) * x_hat + self.beta.view(1,-1,1,1)

小批量数据问题：
当 batch_size 较小时，方差计算可能不稳定
解决方案：增大 eps 值（默认 1e-5），或使用 BatchNorm 的momentum参数调整统计量更新速度
模式切换陷阱：
训练结束忘记调用 model.eval() 会导致推理结果不一致
验证阶段建议使用 with torch.no_grad(): 包裹前向计算
混合精度训练：
BN 层的统计量计算建议保持 FP32 精度
PyTorch 中可通过 torch.cuda.amp.autocast(enabled=False) 包裹 BN 层

# 自定义 BN 与官方实现的误差测试
custom_bn = CustomBatchNorm2d(64)
official_bn = nn.BatchNorm2d(64)

# 参数同步
official_bn.weight.data = custom_bn.gamma.data.clone()
official_bn.bias.data = custom_bn.beta.data.clone()

x = torch.randn(32, 64, 128, 128)
y_custom = custom_bn(x)
y_official = official_bn(x)

torch.allclose(y_custom, y_official, atol=1e-6)  # 应返回 True