Batch Norm反向传播实现详解：从数学推导到PyTorch最佳实践

1次阅读

没有评论

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

Batch Normalization（BN）是现代深度学习模型中不可或缺的组件，它能显著加速训练并提升模型性能。然而，BN 层的反向传播实现却暗藏诸多陷阱：

梯度不稳定：在手动实现时，若未正确处理均值 / 方差的梯度计算，极易引发梯度爆炸或消失。这是因为 BN 涉及对 batch 统计量的依赖，使得梯度流经的路径比普通层更复杂
模式切换隐患：训练 / 推理模式的不当切换会导致统计量更新错误，表现为推理时性能突然下降
数值敏感性：方差计算中的分母可能接近零，若无保护措施会导致 NaN 问题

给定输入 $X \in \mathbb{R}^{N\times C\times H\times W}$（N 为 batch size），BN 的计算分为三步：

计算 batch 统计量：
$$\mu = \frac{1}{NHW}\sum_{n,h,w} X_{n,c,h,w}$$
$$\sigma^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$$} – \mu_c}{\sqrt{\sigma_c^2}
仿射变换：
$$Y_{n,c,h,w} = \gamma_c \hat{X}_{n,c,h,w} + \beta_c$$

设上游梯度为 $\frac{\partial L}{\partial Y}$，需计算三个关键梯度：

参数梯度：
$$\frac{\partial L}{\partial \gamma} = \sum_{n,h,w} \frac{\partial L}{\partial Y_{n,c,h,w}} \hat{X}{n,c,h,w}$$
$$\frac{\partial L}{\partial \beta} = \sum$$} \frac{\partial L}{\partial Y_{n,c,h,w}
输入梯度（推导过程涉及多元链式法则）：
$$\frac{\partial L}{\partial X} = \frac{\gamma}{\sqrt{\sigma^2+\epsilon}} \left(\frac{\partial L}{\partial Y} – \frac{1}{NHW}\left(\sum \frac{\partial L}{\partial Y} + \hat{X} \sum \frac{\partial L}{\partial Y}\hat{X}\right)\right)$$

import torch
import torch.nn as nn

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

        # 注册 buffer 用于推理模式
        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 统计量
            dims = (0, 2, 3)
            mean = x.mean(dims, keepdim=True)
            var = x.var(dims, unbiased=False, keepdim=True)

            # 更新 running stats
            with torch.no_grad():
                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()
        else:
            # 推理模式：使用 running stats
            mean = self.running_mean.view(1, -1, 1, 1)
            var = self.running_var.view(1, -1, 1, 1)

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

关键实现细节：