深度学习入门：从零理解BN层的反向传播实现原理

1次阅读

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

批量归一化（Batch Normalization, BN）是深度学习中常用的技术，主要用于加速神经网络的训练过程并提高模型的稳定性。它的核心思想是对每一层的输入进行归一化处理，使得输入数据的分布更加稳定。然而，对于初学者来说，BN 层的反向传播过程往往是一个难点，特别是涉及到均值、方差的梯度计算时，容易感到困惑。

BN 层的作用 ：BN 层通过对每个 batch 的数据进行归一化，使得输入数据的均值为 0，方差为 1。这样可以减少内部协变量偏移（Internal Covariate Shift），从而加速模型的收敛速度。
常见困惑点 ：初学者在理解 BN 层的反向传播时，常常对如何计算均值、方差的梯度感到困惑。此外，BN 层在训练和推理阶段的处理方式不同，这也容易引发混淆。

BN 层的前向传播和反向传播过程可以用数学公式清晰地表达出来。以下是详细的推导过程：

前向传播 ：
计算 batch 数据的均值：
[\mu = \frac{1}{m} \sum_{i=1}^m x_i ]
计算 batch 数据的方差：
[\sigma^2 = \frac{1}{m} \sum_{i=1}^m (x_i – \mu)^2 ]
对数据进行归一化：
[\hat{x}_i = \frac{x_i – \mu}{\sqrt{\sigma^2 + \epsilon}} ]
应用缩放和偏移参数：
[y_i = \gamma \hat{x}_i + \beta ]
反向传播 ：
计算梯度时，需要分别对输入数据、缩放参数和偏移参数进行梯度计算。
对于输入数据的梯度，可以通过链式法则逐步推导。
缩放参数和偏移参数的梯度计算相对简单，因为它们直接作用于归一化后的数据。

以下是 PyTorch 中 BN 层反向传播的简化实现代码，代码中包含了关键步骤的注释：

import torch
import torch.nn as nn

class BatchNorm1d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super(BatchNorm1d, self).__init__()
        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)
            # 更新 running_mean 和 running_var
            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)
        # 应用缩放和偏移
        y = self.gamma * x_hat + self.beta
        return y

batch size 过小 ：当 batch size 过小时，计算出的均值和方差可能不够准确，从而导致训练不稳定。建议 batch size 不要小于 32。
推理阶段的处理 ：在推理阶段，BN 层使用的是训练过程中累积的 running_mean 和 running_var，而不是当前 batch 的统计量。这一点在实现时需要特别注意。
初始化参数 ：缩放参数（gamma）和偏移参数（beta）的初始化也很重要。通常 gamma 初始化为 1，beta 初始化为 0。