8.5. 批量规范化
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 SageMaker Studio Lab 中打开 Notebook

训练深度神经网络很困难,让它们在合理的时间内收敛可能是一个挑战。在本节中,我们介绍*批量规范化*(batch normalization),这是一种流行且有效的技术,可持续加速深层网络的收敛 (Ioffe and Szegedy, 2015)。再结合在 8.6节 中将介绍的残差块,批量规范化使得研究人员能够训练出层数超过100的。批量规范化的一个(意外)好处在于其固有的正则化效果。

import torch
from torch import nn
from d2l import torch as d2l
from mxnet import autograd, init, np, npx
from mxnet.gluon import nn
from d2l import mxnet as d2l

npx.set_np()
from functools import partial
import jax
import optax
from flax import linen as nn
from jax import numpy as jnp
from d2l import jax as d2l
import tensorflow as tf
from d2l import tensorflow as d2l

8.5.1. 训练深层网络

在处理数据时,我们通常在训练前进行预处理。关于数据预处理的选择往往会对最终结果产生巨大的影响。回想一下我们在 5.7节 中将多层感知机应用于预测房价的例子。在处理真实数据时,我们的第一步是标准化输入特征,使其在多个观测值上具有零均值 \(\boldsymbol{\mu} = 0\) 和单位方差 \(\boldsymbol{\Sigma} = \boldsymbol{1}\) (Friedman, 1987),并经常重新缩放后者,使其对角线为1,即 \(\Sigma_{ii} = 1\)。另一种策略是*对每个观测值*将向量重新缩放到单位长度,可能为零均值。这在空间传感器数据等情况下效果很好。这些预处理技术和其他许多技术都有助于保持估计问题处于良好控制之下。关于特征选择和提取的综述,可参见例如 Guyon等人 (2008) 的文章。标准化向量还有一个很好的副作用,即限制了作用于其上的函数的函数复杂性。例如,支持向量机中著名的半径-边际界 (Vapnik, 1995) 和感知器收敛定理 (Novikoff, 1962) 都依赖于有界范数的输入。

直观地说,这种标准化方法与我们的优化器配合得很好,因为它*先验地*将参数置于相似的尺度上。因此,很自然地会问,在深度网络*内部*进行相应的规范化步骤是否会有益。虽然这并非是发明批量规范化 (Ioffe and Szegedy, 2015) 的原因,但这是一种有用的方式来理解它及其近亲——层规范化 (Ba et al., 2016),并将它们置于一个统一的框架中。

其次,对于一个典型的多层感知机或卷积神经网络,当我们训练时,中间层中的变量(例如,多层感知机中的仿射变换输出)的数值可能会呈现出数量级的巨大差异:无论是从输入到输出的各层之间,还是同一层内的单元之间,以及随着我们对模型参数的更新而随时间变化。批量规范化的发明者非正式地假设,这种变量分布的漂移可能会妨碍网络的收敛。直观地,我们可能会猜想,如果某一层变量的激活值是另一层的100倍,这可能需要学习率进行补偿性调整。像AdaGrad (Duchi et al., 2011)、Adam (Kingma and Ba, 2014)、Yogi (Zaheer et al., 2018)或Distributed Shampoo (Anil et al., 2020) 等自适应求解器旨在从优化的角度解决这个问题,例如通过添加二阶方法的方面。另一种方法是通过自适应规范化来防止问题发生。

第三,更深层的网络是复杂的,并且更容易过拟合。这意味着正则化变得更加关键。正则化的一种常用技术是噪声注入。这早已为人所知,例如,关于输入噪声注入 (Bishop, 1995)。它也构成了 5.6节 中丢弃法的基础。事实证明,非常偶然地,批量规范化同时带来了这三个好处:预处理、数值稳定性和正则化。

批量规范化应用于单个层,或者可选地,应用于所有层:在每次训练迭代中,我们首先通过减去其均值并除以其标准差来规范化(批量规范化的)输入,其中均值和标准差都是基于当前小批量的统计数据估计的。接下来,我们应用一个缩放系数和一个偏移量来恢复失去的自由度。正是由于这种基于*批量*统计数据的*规范化*,*批量规范化*才得此名。

请注意,如果我们尝试使用大小为1的小批量来应用批量规范化,我们将无法学到任何东西。这是因为在减去均值后,每个隐藏单元都会取值为0。正如你可能猜到的,既然我们专门用一整节来讨论批量规范化,那么只要小批量足够大,这种方法就是有效且稳定的。这里的一个要点是,在应用批量规范化时,批量大小的选择比不使用批量规范化时更为重要,或者至少,在我们调整批量大小时需要进行适当的校准。

\(\mathcal{B}\) 表示一个小批量,令 \(\mathbf{x} \in \mathcal{B}\) 是批量规范化(\(\textrm{BN}\))的一个输入。在这种情况下,批量规范化定义如下

(8.5.1)\[\textrm{BN}(\mathbf{x}) = \boldsymbol{\gamma} \odot \frac{\mathbf{x} - \hat{\boldsymbol{\mu}}_\mathcal{B}}{\hat{\boldsymbol{\sigma}}_\mathcal{B}} + \boldsymbol{\beta}.\]

(8.5.1) 中,\(\hat{\boldsymbol{\mu}}_\mathcal{B}\) 是样本均值,\(\hat{\boldsymbol{\sigma}}_\mathcal{B}\) 是小批量 \(\mathcal{B}\) 的样本标准差。应用标准化后,得到的小批量具有零均值和单位方差。选择单位方差(而不是其他某个神奇的数字)是任意的。我们通过包含一个逐元素的*缩放参数* \(\boldsymbol{\gamma}\) 和*平移参数* \(\boldsymbol{\beta}\) 来恢复这个自由度,它们的形状与 \(\mathbf{x}\) 相同。这两个参数都需要作为模型训练的一部分来学习。

中间层的变量量级在训练期间不会发散,因为批量规范化通过 \(\hat{\boldsymbol{\mu}}_\mathcal{B}\)\({\hat{\boldsymbol{\sigma}}_\mathcal{B}}\) 主动地将它们居中并重新缩放到给定的均值和大小。实际经验证实,正如在讨论特征重新缩放时所提到的,批量规范化似乎允许更激进的学习率。我们在 (8.5.1) 中计算 \(\hat{\boldsymbol{\mu}}_\mathcal{B}\)\({\hat{\boldsymbol{\sigma}}_\mathcal{B}}\) 如下

(8.5.2)\[\hat{\boldsymbol{\mu}}_\mathcal{B} = \frac{1}{|\mathcal{B}|} \sum_{\mathbf{x} \in \mathcal{B}} \mathbf{x} \textrm{ 和 } \hat{\boldsymbol{\sigma}}_\mathcal{B}^2 = \frac{1}{|\mathcal{B}|} \sum_{\mathbf{x} \in \mathcal{B}} (\mathbf{x} - \hat{\boldsymbol{\mu}}_{\mathcal{B}})^2 + \epsilon.\]

注意,我们在方差估计中添加了一个小的常数 \(\epsilon > 0\),以确保我们永远不会尝试除以零,即使在经验方差估计可能非常小或消失的情况下也是如此。估计值 \(\hat{\boldsymbol{\mu}}_\mathcal{B}\)\({\hat{\boldsymbol{\sigma}}_\mathcal{B}}\) 通过使用均值和方差的噪声估计来抵消缩放问题。你可能会认为这种噪声应该是一个问题。恰恰相反,它实际上是有益的。

这在深度学习中是一个反复出现的主题。出于尚未在理论上得到很好表征的原因,优化中的各种噪声源通常会导致更快的训练和更少的过拟合:这种变化似乎起到了一种正则化的作用。Teye等人 (2018)Luo等人 (2018) 分别将批量规范化的性质与贝叶斯先验和惩罚联系起来。特别是,这为为什么批量规范化在50-100范围内的中等小批量大小下效果最好这一难题提供了一些启示。这种特定大小的小批量似乎为每一层注入了“恰到好处”的噪声,无论是通过 \(\hat{\boldsymbol{\sigma}}\) 在尺度方面,还是通过 \(\hat{\boldsymbol{\mu}}\) 在偏移方面:较大的小批量由于估计更稳定而正则化效果较差,而极小的小批量由于高方差而破坏了有用的信号。进一步探索这个方向,考虑其他类型的预处理和滤波,可能会产生其他有效的正则化类型。

固定一个训练好的模型,你可能会认为我们更愿意使用整个数据集来估计均值和方差。一旦训练完成,我们为什么会希望同一张图像根据它恰好所在的批次而被不同地分类呢?在训练期间,这种精确计算是不可行的,因为每次我们更新模型时,所有数据示例的中间变量都会改变。然而,一旦模型训练完成,我们就可以根据整个数据集计算每一层变量的均值和方差。实际上,这对于采用批量规范化的模型来说是标准做法;因此,批量规范化层在*训练模式*(通过小批量统计数据进行规范化)和*预测模式*(通过数据集统计数据进行规范化)下的功能是不同的。在这种形式下,它们的行为与 5.6节 中的丢弃正则化非常相似,其中噪声仅在训练期间注入。

8.5.2. 批量规范化层

全连接层和卷积层的批量规范化实现略有不同。批量规范化与其他层的一个关键区别在于,因为前者一次对整个小批量进行操作,我们不能像之前介绍其他层时那样忽略批量维度。

8.5.2.1. 全连接层

当将批量规范化应用于全连接层时,Ioffe和Szegedy (2015) 在他们的原始论文中将批量规范化插入到仿射变换之后和非线性激活函数*之前*。后来的应用尝试将批量规范化插入到激活函数*之后*。用 \(\mathbf{x}\) 表示全连接层的输入,仿射变换为 \(\mathbf{W}\mathbf{x} + \mathbf{b}\)(其中权重参数为 \(\mathbf{W}\),偏置参数为 \(\mathbf{b}\)),激活函数为 \(\phi\),我们可以将启用了批量规范化的全连接层输出 \(\mathbf{h}\) 的计算表示如下:

(8.5.3)\[\mathbf{h} = \phi(\textrm{BN}(\mathbf{W}\mathbf{x} + \mathbf{b}) ).\]

回想一下,均值和方差是在应用变换的*同一个*小批量上计算的。

8.5.2.2. 卷积层

同样,对于卷积层,我们可以在卷积之后、非线性激活函数之前应用批量规范化。与全连接层中的批量规范化的关键区别在于,我们是*跨所有位置*对每个通道应用该操作。这与我们导致卷积的平移不变性假设是一致的:我们假设图像中模式的具体位置对于理解的目的并不重要。

假设我们的小批量包含 \(m\) 个样本,并且对于每个通道,卷积的输出高度为 \(p\),宽度为 \(q\)。对于卷积层,我们同时对每个输出通道的 \(m \cdot p \cdot q\) 个元素进行批量规范化。因此,我们在计算均值和方差时收集所有空间位置的值,并因此在给定通道内应用相同的均值和方差来规范化每个空间位置的值。每个通道都有自己的缩放和移位参数,它们都是标量。

8.5.2.3. 层规范化

请注意,在卷积的背景下,批量规范化即使对于大小为1的小批量也是明确定义的:毕竟,我们有图像上的所有位置可以取平均。因此,即使只是在单个观测内,均值和方差也是明确定义的。这个考虑导致了Ba等人 (2016) 引入了*层规范化*的概念。它的工作方式就像批量规范化一样,只是它一次应用于一个观测。因此,偏移和缩放因子都是标量。对于一个 \(n\) 维向量 \(\mathbf{x}\),层规范化由下式给出:

(8.5.4)\[\mathbf{x} \rightarrow \textrm{LN}(\mathbf{x}) = \frac{\mathbf{x} - \hat{\mu}}{\hat\sigma},\]

其中缩放和偏移是逐系数应用的,并由下式给出:

(8.5.5)\[\hat{\mu} \stackrel{\textrm{def}}{=} \frac{1}{n} \sum_{i=1}^n x_i \textrm{ and } \hat{\sigma}^2 \stackrel{\textrm{def}}{=} \frac{1}{n} \sum_{i=1}^n (x_i - \hat{\mu})^2 + \epsilon.\]

和以前一样,我们添加了一个小的偏移量 \(\epsilon > 0\) 以防止除以零。使用层规范化的一个主要好处是它可以防止发散。毕竟,忽略 \(\epsilon\),层规范化的输出是尺度无关的。也就是说,对于任何 \(\alpha \neq 0\) 的选择,我们有 \(\textrm{LN}(\mathbf{x}) \approx \textrm{LN}(\alpha \mathbf{x})\)。当 \(|\alpha| \to \infty\) 时,这成为一个等式(近似等式是由于方差的偏移量 \(\epsilon\))。

层规范化的另一个优点是它不依赖于小批量的大小。它也与我们是在训练模式还是测试模式无关。换句话说,它只是一个将激活值标准化到给定尺度的确定性变换。这在防止优化中的发散方面非常有益。我们跳过更多细节,建议感兴趣的读者查阅原始论文。

8.5.2.4. 预测期间的批量规范化

正如我们前面提到的,批量规范化在训练模式和预测模式下的行为通常不同。首先,一旦我们训练好模型,从小批量上估计样本均值和样本方差所产生的噪声就不再是理想的了。其次,我们可能没有计算每批规范化统计数据的奢侈。例如,我们可能需要应用我们的模型一次只做一个预测。

通常,在训练之后,我们使用整个数据集来计算变量统计数据的稳定估计,然后在预测时固定它们。因此,批量规范化在训练和测试时的行为是不同的。回想一下,丢弃法也表现出这种特性。

8.5.3. 从零开始实现

为了了解批量规范化在实践中是如何工作的,我们下面从头开始实现一个。

def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
    # Use is_grad_enabled to determine whether we are in training mode
    if not torch.is_grad_enabled():
        # In prediction mode, use mean and variance obtained by moving average
        X_hat = (X - moving_mean) / torch.sqrt(moving_var + eps)
    else:
        assert len(X.shape) in (2, 4)
        if len(X.shape) == 2:
            # When using a fully connected layer, calculate the mean and
            # variance on the feature dimension
            mean = X.mean(dim=0)
            var = ((X - mean) ** 2).mean(dim=0)
        else:
            # When using a two-dimensional convolutional layer, calculate the
            # mean and variance on the channel dimension (axis=1). Here we
            # need to maintain the shape of X, so that the broadcasting
            # operation can be carried out later
            mean = X.mean(dim=(0, 2, 3), keepdim=True)
            var = ((X - mean) ** 2).mean(dim=(0, 2, 3), keepdim=True)
        # In training mode, the current mean and variance are used
        X_hat = (X - mean) / torch.sqrt(var + eps)
        # Update the mean and variance using moving average
        moving_mean = (1.0 - momentum) * moving_mean + momentum * mean
        moving_var = (1.0 - momentum) * moving_var + momentum * var
    Y = gamma * X_hat + beta  # Scale and shift
    return Y, moving_mean.data, moving_var.data
def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
    # Use autograd to determine whether we are in training mode
    if not autograd.is_training():
        # In prediction mode, use mean and variance obtained by moving average
        X_hat = (X - moving_mean) / np.sqrt(moving_var + eps)
    else:
        assert len(X.shape) in (2, 4)
        if len(X.shape) == 2:
            # When using a fully connected layer, calculate the mean and
            # variance on the feature dimension
            mean = X.mean(axis=0)
            var = ((X - mean) ** 2).mean(axis=0)
        else:
            # When using a two-dimensional convolutional layer, calculate the
            # mean and variance on the channel dimension (axis=1). Here we
            # need to maintain the shape of X, so that the broadcasting
            # operation can be carried out later
            mean = X.mean(axis=(0, 2, 3), keepdims=True)
            var = ((X - mean) ** 2).mean(axis=(0, 2, 3), keepdims=True)
        # In training mode, the current mean and variance are used
        X_hat = (X - mean) / np.sqrt(var + eps)
        # Update the mean and variance using moving average
        moving_mean = (1.0 - momentum) * moving_mean + momentum * mean
        moving_var = (1.0 - momentum) * moving_var + momentum * var
    Y = gamma * X_hat + beta  # Scale and shift
    return Y, moving_mean, moving_var
def batch_norm(X, deterministic, gamma, beta, moving_mean, moving_var, eps,
               momentum):
    # Use `deterministic` to determine whether the current mode is training
    # mode or prediction mode
    if deterministic:
        # In prediction mode, use mean and variance obtained by moving average
        # `linen.Module.variables` have a `value` attribute containing the array
        X_hat = (X - moving_mean.value) / jnp.sqrt(moving_var.value + eps)
    else:
        assert len(X.shape) in (2, 4)
        if len(X.shape) == 2:
            # When using a fully connected layer, calculate the mean and
            # variance on the feature dimension
            mean = X.mean(axis=0)
            var = ((X - mean) ** 2).mean(axis=0)
        else:
            # When using a two-dimensional convolutional layer, calculate the
            # mean and variance on the channel dimension (axis=1). Here we
            # need to maintain the shape of `X`, so that the broadcasting
            # operation can be carried out later
            mean = X.mean(axis=(0, 2, 3), keepdims=True)
            var = ((X - mean) ** 2).mean(axis=(0, 2, 3), keepdims=True)
        # In training mode, the current mean and variance are used
        X_hat = (X - mean) / jnp.sqrt(var + eps)
        # Update the mean and variance using moving average
        moving_mean.value = momentum * moving_mean.value + (1.0 - momentum) * mean
        moving_var.value = momentum * moving_var.value + (1.0 - momentum) * var
    Y = gamma * X_hat + beta  # Scale and shift
    return Y
def batch_norm(X, gamma, beta, moving_mean, moving_var, eps):
    # Compute reciprocal of square root of the moving variance elementwise
    inv = tf.cast(tf.math.rsqrt(moving_var + eps), X.dtype)
    # Scale and shift
    inv *= gamma
    Y = X * inv + (beta - moving_mean * inv)
    return Y

我们现在可以创建一个合适的 BatchNorm 层。我们的层将为缩放 gamma 和平移 beta 维护适当的参数,这两者都将在训练过程中更新。此外,我们的层将维护均值和方差的移动平均值,以便在模型预测期间后续使用。

撇开算法细节不谈,请注意我们实现该层所遵循的设计模式。通常,我们会将数学运算定义在一个单独的函数中,例如 batch_norm。然后,我们将此功能集成到一个自定义层中,其代码主要处理簿记事宜,例如将数据移动到正确的设备上下文,分配和初始化任何必需的变量,跟踪移动平均值(此处为均值和方差)等。这种模式能够将数学与样板代码清晰地分离开来。另请注意,为方便起见,我们在这里没有考虑自动推断输入形状;因此我们需要在整个过程中指定特征数量。到现在为止,所有现代深度学习框架都在高级批量规范化API中提供了大小和形状的自动检测(在实践中我们将使用这个)。

class BatchNorm(nn.Module):
    # num_features: the number of outputs for a fully connected layer or the
    # number of output channels for a convolutional layer. num_dims: 2 for a
    # fully connected layer and 4 for a convolutional layer
    def __init__(self, num_features, num_dims):
        super().__init__()
        if num_dims == 2:
            shape = (1, num_features)
        else:
            shape = (1, num_features, 1, 1)
        # The scale parameter and the shift parameter (model parameters) are
        # initialized to 1 and 0, respectively
        self.gamma = nn.Parameter(torch.ones(shape))
        self.beta = nn.Parameter(torch.zeros(shape))
        # The variables that are not model parameters are initialized to 0 and
        # 1
        self.moving_mean = torch.zeros(shape)
        self.moving_var = torch.ones(shape)

    def forward(self, X):
        # If X is not on the main memory, copy moving_mean and moving_var to
        # the device where X is located
        if self.moving_mean.device != X.device:
            self.moving_mean = self.moving_mean.to(X.device)
            self.moving_var = self.moving_var.to(X.device)
        # Save the updated moving_mean and moving_var
        Y, self.moving_mean, self.moving_var = batch_norm(
            X, self.gamma, self.beta, self.moving_mean,
            self.moving_var, eps=1e-5, momentum=0.1)
        return Y
class BatchNorm(nn.Block):
    # `num_features`: the number of outputs for a fully connected layer
    # or the number of output channels for a convolutional layer. `num_dims`:
    # 2 for a fully connected layer and 4 for a convolutional layer
    def __init__(self, num_features, num_dims, **kwargs):
        super().__init__(**kwargs)
        if num_dims == 2:
            shape = (1, num_features)
        else:
            shape = (1, num_features, 1, 1)
        # The scale parameter and the shift parameter (model parameters) are
        # initialized to 1 and 0, respectively
        self.gamma = self.params.get('gamma', shape=shape, init=init.One())
        self.beta = self.params.get('beta', shape=shape, init=init.Zero())
        # The variables that are not model parameters are initialized to 0 and
        # 1
        self.moving_mean = np.zeros(shape)
        self.moving_var = np.ones(shape)

    def forward(self, X):
        # If `X` is not on the main memory, copy `moving_mean` and
        # `moving_var` to the device where `X` is located
        if self.moving_mean.ctx != X.ctx:
            self.moving_mean = self.moving_mean.copyto(X.ctx)
            self.moving_var = self.moving_var.copyto(X.ctx)
        # Save the updated `moving_mean` and `moving_var`
        Y, self.moving_mean, self.moving_var = batch_norm(
            X, self.gamma.data(), self.beta.data(), self.moving_mean,
            self.moving_var, eps=1e-12, momentum=0.1)
        return Y
class BatchNorm(nn.Module):
    # `num_features`: the number of outputs for a fully connected layer
    # or the number of output channels for a convolutional layer.
    # `num_dims`: 2 for a fully connected layer and 4 for a convolutional layer
    # Use `deterministic` to determine whether the current mode is training
    # mode or prediction mode
    num_features: int
    num_dims: int
    deterministic: bool = False

    @nn.compact
    def __call__(self, X):
        if self.num_dims == 2:
            shape = (1, self.num_features)
        else:
            shape = (1, 1, 1, self.num_features)

        # The scale parameter and the shift parameter (model parameters) are
        # initialized to 1 and 0, respectively
        gamma = self.param('gamma', jax.nn.initializers.ones, shape)
        beta = self.param('beta', jax.nn.initializers.zeros, shape)

        # The variables that are not model parameters are initialized to 0 and
        # 1. Save them to the 'batch_stats' collection
        moving_mean = self.variable('batch_stats', 'moving_mean', jnp.zeros, shape)
        moving_var = self.variable('batch_stats', 'moving_var', jnp.ones, shape)
        Y = batch_norm(X, self.deterministic, gamma, beta,
                       moving_mean, moving_var, eps=1e-5, momentum=0.9)

        return Y
class BatchNorm(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(BatchNorm, self).__init__(**kwargs)

    def build(self, input_shape):
        weight_shape = [input_shape[-1], ]
        # The scale parameter and the shift parameter (model parameters) are
        # initialized to 1 and 0, respectively
        self.gamma = self.add_weight(name='gamma', shape=weight_shape,
            initializer=tf.initializers.ones, trainable=True)
        self.beta = self.add_weight(name='beta', shape=weight_shape,
            initializer=tf.initializers.zeros, trainable=True)
        # The variables that are not model parameters are initialized to 0
        self.moving_mean = self.add_weight(name='moving_mean',
            shape=weight_shape, initializer=tf.initializers.zeros,
            trainable=False)
        self.moving_variance = self.add_weight(name='moving_variance',
            shape=weight_shape, initializer=tf.initializers.ones,
            trainable=False)
        super(BatchNorm, self).build(input_shape)

    def assign_moving_average(self, variable, value):
        momentum = 0.1
        delta = (1.0 - momentum) * variable + momentum * value
        return variable.assign(delta)

    @tf.function
    def call(self, inputs, training):
        if training:
            axes = list(range(len(inputs.shape) - 1))
            batch_mean = tf.reduce_mean(inputs, axes, keepdims=True)
            batch_variance = tf.reduce_mean(tf.math.squared_difference(
                inputs, tf.stop_gradient(batch_mean)), axes, keepdims=True)
            batch_mean = tf.squeeze(batch_mean, axes)
            batch_variance = tf.squeeze(batch_variance, axes)
            mean_update = self.assign_moving_average(
                self.moving_mean, batch_mean)
            variance_update = self.assign_moving_average(
                self.moving_variance, batch_variance)
            self.add_update(mean_update)
            self.add_update(variance_update)
            mean, variance = batch_mean, batch_variance
        else:
            mean, variance = self.moving_mean, self.moving_variance
        output = batch_norm(inputs, moving_mean=mean, moving_var=variance,
            beta=self.beta, gamma=self.gamma, eps=1e-5)
        return output

我们使用 momentum 来控制过去均值和方差估计的聚合。这有点用词不当,因为它与优化的*动量*项毫无关系。尽管如此,这是这个术语的通用名称,为了尊重API命名惯例,我们在代码中使用了相同的变量名。

8.5.4. 使用批量规范化的LeNet

为了了解如何在上下文中应用 BatchNorm,下面我们将其应用于传统的 LeNet 模型(7.6节)。回想一下,批量规范化应用于卷积层或全连接层之后,但在相应的激活函数之前。

class BNLeNetScratch(d2l.Classifier):
    def __init__(self, lr=0.1, num_classes=10):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Sequential(
            nn.LazyConv2d(6, kernel_size=5), BatchNorm(6, num_dims=4),
            nn.Sigmoid(), nn.AvgPool2d(kernel_size=2, stride=2),
            nn.LazyConv2d(16, kernel_size=5), BatchNorm(16, num_dims=4),
            nn.Sigmoid(), nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Flatten(), nn.LazyLinear(120),
            BatchNorm(120, num_dims=2), nn.Sigmoid(), nn.LazyLinear(84),
            BatchNorm(84, num_dims=2), nn.Sigmoid(),
            nn.LazyLinear(num_classes))
class BNLeNetScratch(d2l.Classifier):
    def __init__(self, lr=0.1, num_classes=10):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Sequential()
        self.net.add(
            nn.Conv2D(6, kernel_size=5), BatchNorm(6, num_dims=4),
            nn.Activation('sigmoid'),
            nn.AvgPool2D(pool_size=2, strides=2),
            nn.Conv2D(16, kernel_size=5), BatchNorm(16, num_dims=4),
            nn.Activation('sigmoid'),
            nn.AvgPool2D(pool_size=2, strides=2), nn.Dense(120),
            BatchNorm(120, num_dims=2), nn.Activation('sigmoid'),
            nn.Dense(84), BatchNorm(84, num_dims=2),
            nn.Activation('sigmoid'), nn.Dense(num_classes))
        self.initialize()
class BNLeNetScratch(d2l.Classifier):
    lr: float = 0.1
    num_classes: int = 10
    training: bool = True

    def setup(self):
        self.net = nn.Sequential([
            nn.Conv(6, kernel_size=(5, 5)),
            BatchNorm(6, num_dims=4, deterministic=not self.training),
            nn.sigmoid,
            lambda x: nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2)),
            nn.Conv(16, kernel_size=(5, 5)),
            BatchNorm(16, num_dims=4, deterministic=not self.training),
            nn.sigmoid,
            lambda x: nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2)),
            lambda x: x.reshape((x.shape[0], -1)),
            nn.Dense(120),
            BatchNorm(120, num_dims=2, deterministic=not self.training),
            nn.sigmoid,
            nn.Dense(84),
            BatchNorm(84, num_dims=2, deterministic=not self.training),
            nn.sigmoid,
            nn.Dense(self.num_classes)])

由于 BatchNorm 层需要计算批量统计数据(均值和方差),Flax 会跟踪 batch_stats 字典,并随着每个小批量更新它们。像 batch_stats 这样的集合可以存储在 TrainState 对象中(在 3.2.4节 中定义的 d2l.Trainer 类中)作为一个属性,在模型的前向传播期间,这些应该传递给 mutable 参数,这样 Flax 就会返回突变的变量。

@d2l.add_to_class(d2l.Classifier)  #@save
@partial(jax.jit, static_argnums=(0, 5))
def loss(self, params, X, Y, state, averaged=True):
    Y_hat, updates = state.apply_fn({'params': params,
                                     'batch_stats': state.batch_stats},
                                    *X, mutable=['batch_stats'],
                                    rngs={'dropout': state.dropout_rng})
    Y_hat = Y_hat.reshape((-1, Y_hat.shape[-1]))
    Y = Y.reshape((-1,))
    fn = optax.softmax_cross_entropy_with_integer_labels
    return (fn(Y_hat, Y).mean(), updates) if averaged else (fn(Y_hat, Y), updates)
class BNLeNetScratch(d2l.Classifier):
    def __init__(self, lr=0.1, num_classes=10):
        super().__init__()
        self.save_hyperparameters()
        self.net = tf.keras.models.Sequential([
            tf.keras.layers.Conv2D(filters=6, kernel_size=5,
                                   input_shape=(28, 28, 1)),
            BatchNorm(), tf.keras.layers.Activation('sigmoid'),
            tf.keras.layers.AvgPool2D(pool_size=2, strides=2),
            tf.keras.layers.Conv2D(filters=16, kernel_size=5),
            BatchNorm(), tf.keras.layers.Activation('sigmoid'),
            tf.keras.layers.AvgPool2D(pool_size=2, strides=2),
            tf.keras.layers.Flatten(), tf.keras.layers.Dense(120),
            BatchNorm(), tf.keras.layers.Activation('sigmoid'),
            tf.keras.layers.Dense(84), BatchNorm(),
            tf.keras.layers.Activation('sigmoid'),
            tf.keras.layers.Dense(num_classes)])

和以前一样,我们将在 Fashion-MNIST 数据集上训练我们的网络。这段代码与我们首次训练 LeNet 时的代码几乎完全相同。

trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
data = d2l.FashionMNIST(batch_size=128)
model = BNLeNetScratch(lr=0.1)
model.apply_init([next(iter(data.get_dataloader(True)))[0]], d2l.init_cnn)
trainer.fit(model, data)
../_images/output_batch-norm_cf033c_65_0.svg
trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
data = d2l.FashionMNIST(batch_size=128)
model = BNLeNetScratch(lr=0.1)
trainer.fit(model, data)
../_images/output_batch-norm_cf033c_68_0.svg
trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
data = d2l.FashionMNIST(batch_size=128)
model = BNLeNetScratch(lr=0.1)
trainer.fit(model, data)
../_images/output_batch-norm_cf033c_71_0.svg
trainer = d2l.Trainer(max_epochs=10)
data = d2l.FashionMNIST(batch_size=128)
with d2l.try_gpu():
    model = BNLeNetScratch(lr=0.5)
    trainer.fit(model, data)
../_images/output_batch-norm_cf033c_74_0.svg

让我们看一看从第一个批量规范化层学到的缩放参数 gamma 和平移参数 beta

model.net[1].gamma.reshape((-1,)), model.net[1].beta.reshape((-1,))
(tensor([1.4334, 1.9905, 1.8584, 2.0740, 2.0522, 1.8877], device='cuda:0',
        grad_fn=<ViewBackward0>),
 tensor([ 0.7354, -1.3538, -0.2567, -0.9991, -0.3028,  1.3125], device='cuda:0',
        grad_fn=<ViewBackward0>))
model.net[1].gamma.data().reshape(-1,), model.net[1].beta.data().reshape(-1,)
(array([2.130113 , 1.560813 , 1.461431 , 1.9807949, 2.2318861, 1.551563 ], ctx=gpu(0)),
 array([ 1.2277379 ,  1.514598  , -1.014917  ,  0.19028394,  0.6355166 ,
         0.5642359 ], ctx=gpu(0)))
trainer.state.params['net']['layers_1']['gamma'].reshape((-1,)), \
trainer.state.params['net']['layers_1']['beta'].reshape((-1,))
(Array([1.6042372, 2.336828 , 1.6758277, 2.621987 , 1.1613046, 2.3155787],      dtype=float32),
 Array([-0.21946815,  0.64822614, -0.3630768 ,  0.88309413, -0.04114898,
        -0.610043  ], dtype=float32))
tf.reshape(model.net.layers[1].gamma, (-1,)), tf.reshape(
    model.net.layers[1].beta, (-1,))
(<tf.Tensor: shape=(6,), dtype=float32, numpy=
 array([3.6447892, 2.797243 , 1.5961126, 1.919072 , 4.717426 , 2.2385437],
       dtype=float32)>,
 <tf.Tensor: shape=(6,), dtype=float32, numpy=
 array([-0.18158795, -0.27821252, -0.83241636, -1.0528361 ,  0.2683058 ,
        -0.6192782 ], dtype=float32)>)

8.5.5. 简洁实现

与我们刚才自己定义的 BatchNorm 类相比,我们可以直接使用深度学习框架高级API中定义的 BatchNorm 类。代码看起来与我们上面的实现几乎完全相同,只是我们不再需要提供额外的参数来使其维度正确。

class BNLeNet(d2l.Classifier):
    def __init__(self, lr=0.1, num_classes=10):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Sequential(
            nn.LazyConv2d(6, kernel_size=5), nn.LazyBatchNorm2d(),
            nn.Sigmoid(), nn.AvgPool2d(kernel_size=2, stride=2),
            nn.LazyConv2d(16, kernel_size=5), nn.LazyBatchNorm2d(),
            nn.Sigmoid(), nn.AvgPool2d(kernel_size=2, stride=2),
            nn.Flatten(), nn.LazyLinear(120), nn.LazyBatchNorm1d(),
            nn.Sigmoid(), nn.LazyLinear(84), nn.LazyBatchNorm1d(),
            nn.Sigmoid(), nn.LazyLinear(num_classes))
class BNLeNet(d2l.Classifier):
    def __init__(self, lr=0.1, num_classes=10):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Sequential()
        self.net.add(
            nn.Conv2D(6, kernel_size=5), nn.BatchNorm(),
            nn.Activation('sigmoid'),
            nn.AvgPool2D(pool_size=2, strides=2),
            nn.Conv2D(16, kernel_size=5), nn.BatchNorm(),
            nn.Activation('sigmoid'),
            nn.AvgPool2D(pool_size=2, strides=2),
            nn.Dense(120), nn.BatchNorm(), nn.Activation('sigmoid'),
            nn.Dense(84), nn.BatchNorm(), nn.Activation('sigmoid'),
            nn.Dense(num_classes))
        self.initialize()
class BNLeNet(d2l.Classifier):
    lr: float = 0.1
    num_classes: int = 10
    training: bool = True

    def setup(self):
        self.net = nn.Sequential([
            nn.Conv(6, kernel_size=(5, 5)),
            nn.BatchNorm(not self.training),
            nn.sigmoid,
            lambda x: nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2)),
            nn.Conv(16, kernel_size=(5, 5)),
            nn.BatchNorm(not self.training),
            nn.sigmoid,
            lambda x: nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2)),
            lambda x: x.reshape((x.shape[0], -1)),
            nn.Dense(120),
            nn.BatchNorm(not self.training),
            nn.sigmoid,
            nn.Dense(84),
            nn.BatchNorm(not self.training),
            nn.sigmoid,
            nn.Dense(self.num_classes)])
class BNLeNet(d2l.Classifier):
    def __init__(self, lr=0.1, num_classes=10):
        super().__init__()
        self.save_hyperparameters()
        self.net = tf.keras.models.Sequential([
            tf.keras.layers.Conv2D(filters=6, kernel_size=5,
                                   input_shape=(28, 28, 1)),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Activation('sigmoid'),
            tf.keras.layers.AvgPool2D(pool_size=2, strides=2),
            tf.keras.layers.Conv2D(filters=16, kernel_size=5),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Activation('sigmoid'),
            tf.keras.layers.AvgPool2D(pool_size=2, strides=2),
            tf.keras.layers.Flatten(), tf.keras.layers.Dense(120),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Activation('sigmoid'),
            tf.keras.layers.Dense(84),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Activation('sigmoid'),
            tf.keras.layers.Dense(num_classes)])

下面,我们使用相同的超参数来训练我们的模型。请注意,像往常一样,高级API变体的运行速度要快得多,因为它的代码已经编译成C++或CUDA,而我们的自定义实现必须由Python解释。

trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
data = d2l.FashionMNIST(batch_size=128)
model = BNLeNet(lr=0.1)
model.apply_init([next(iter(data.get_dataloader(True)))[0]], d2l.init_cnn)
trainer.fit(model, data)
../_images/output_batch-norm_cf033c_110_0.svg
trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
data = d2l.FashionMNIST(batch_size=128)
model = BNLeNet(lr=0.1)
trainer.fit(model, data)
../_images/output_batch-norm_cf033c_113_0.svg
trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
data = d2l.FashionMNIST(batch_size=128)
model = BNLeNet(lr=0.1)
trainer.fit(model, data)
../_images/output_batch-norm_cf033c_116_0.svg
trainer = d2l.Trainer(max_epochs=10)
data = d2l.FashionMNIST(batch_size=128)
with d2l.try_gpu():
    model = BNLeNet(lr=0.5)
    trainer.fit(model, data)
../_images/output_batch-norm_cf033c_119_0.svg

8.5.6. 讨论

直观上,批量规范化被认为可以使优化景观更平滑。然而,我们必须小心区分推测性的直觉和我们训练深度模型时观察到的现象的真实解释。回想一下,我们甚至不知道为什么更简单的深度神经网络(多层感知机和传统卷积神经网络)首先能很好地泛化。即使有丢弃法和权重衰减,它们仍然非常灵活,以至于它们泛化到未见数据的能力可能需要更精细的学习理论泛化保证。

提出批量规范化的原始论文 (Ioffe and Szegedy, 2015),除了介绍一种强大而有用的工具外,还解释了它为什么有效:通过减少*内部协变量偏移*。据推测,他们所说的*内部协变量偏移*指的是上面表达的直觉——即变量值的分布在训练过程中会发生变化。然而,这个解释存在两个问题:i) 这种漂移与*协变量偏移*非常不同,使得这个名称成为一个用词不当。如果说有什么相似之处,它更接近概念漂移。ii) 这个解释提供了一个不明确的直觉,但留下了*为什么这种技术究竟有效*的问题,这是一个悬而未决的问题,需要一个严谨的解释。在本书中,我们旨在传达从业者用来指导他们开发深度神经网络的直觉。然而,我们认为将这些指导性的直觉与既定的科学事实分开是很重要的。最终,当你掌握了这些材料并开始撰写自己的研究论文时,你会希望清楚地划分技术声明和直觉之间的界限。

在批量规范化取得成功之后,其关于*内部协变量偏移*的解释在技术文献和关于如何呈现机器学习研究的更广泛讨论中反复出现。在2017年NeurIPS会议上接受“经受时间考验奖”(Test of Time Award)时发表的一篇令人难忘的演讲中,Ali Rahimi将*内部协变量偏移*作为一个焦点,将现代深度学习实践比作炼金术。随后,这个例子在一篇概述机器学习中令人担忧的趋势的立场文件中被详细回顾 (Lipton and Steinhardt, 2018)。其他作者也为批量规范化的成功提出了替代解释,一些人 (Santurkar et al., 2018) 声称批量规范化的成功是在其表现出与原始论文中声称的行为在某些方面相反的情况下取得的。

我们注意到,*内部协变量偏移*并不比每年在技术机器学习文献中成千上万个类似模糊的主张更值得批评。很可能,它作为这些辩论的焦点所产生的共鸣,归功于其在目标受众中的广泛认知度。批量规范化已被证明是一种不可或缺的方法,应用于几乎所有已部署的图像分类器中,为引入该技术的论文赢得了数万次引用。我们推测,通过噪声注入进行正则化、通过重新缩放加速以及最后的预处理等指导原则,很可能在未来会导致更多层和技术的发明。

从更实际的角度来看,关于批量规范化有几个方面值得记住:

  • 在模型训练过程中,批量规范化利用小批量的均值和标准差不断调整网络的中间输出,使得整个神经网络中每一层的中间输出值更加稳定。

  • 批量规范化对于全连接层和卷积层的处理略有不同。事实上,对于卷积层,层规范化有时可以作为一种替代方案。

  • 与丢弃层一样,批量规范化层在训练模式和预测模式下的行为不同。

  • 批量规范化有助于正则化和改善优化中的收敛性。相比之下,减少内部协变量偏移的最初动机似乎不是一个有效的解释。

  • 对于对输入扰动不太敏感的更鲁棒的模型,可以考虑移除批量规范化 (Wang et al., 2022)

8.5.7. 练习

  1. 在批量规范化之前,我们应该从全连接层或卷积层中移除偏置参数吗?为什么?

  2. 比较使用和不使用批量规范化的LeNet的学习率。

    1. 绘制验证准确率的增长情况。

    2. 在两种情况下,在优化失败之前,你能将学习率设置到多大?

  3. 我们需要在每一层都使用批量规范化吗?试验一下。

  4. 实现一个“精简版”的批量规范化,只移除均值,或者只移除方差。它的表现如何?

  5. 固定参数 betagamma。观察并分析结果。

  6. 你能用批量规范化替代丢弃法吗?行为有何变化?

  7. 研究思路:思考一下你可以应用的其他规范化变换。

    1. 你能应用概率积分变换吗?

    2. 你能使用全秩协方差估计吗?为什么你可能不应该这样做?

    3. 你能使用其他紧凑矩阵变体(块对角、低位移秩、Monarch等)吗?

    4. 稀疏化压缩是否能起到正则化的作用?

    5. 你还能使用其他投影(例如,凸锥、对称群特定变换)吗?