深度学习之卷积神经网络（8）BatchNorm层

卷积神经网络的出现，网络参数量大大减低，使得几十层的深层网络称为可能。然而，在残差网络出现之前，网络的加深使得网络训练变得十分不稳定，甚至出现网络长时间不更新甚至不收敛的现象，同时网络对超参数比较敏感，超参数的微量扰动也会导致网络的训练轨迹完全改变。

 2015年，Google研究人员Sergey Ioffe等提出了一种参数标准化（Normalize）的手段，并基于参数标准化设计了Batch Normalization（简写为BatchNorm，或BN）层。BN层的提出，使得网络的超参数的设定更加自由，比如更大的学习率、更随意的网络初始化等，同时网络的收敛速度更快，性能也更好。BN层提出后便广泛地应用在各种深度网络模型上，卷积层、BN层、ReLU层、池化层一度成为网络模型的标配单元块，通过堆叠Conv-BN-ReLU-Pooling方式往往可以获得不错的模型性能。

BatchNorm层概念

 首先我们来探索，为什么需要对网络中的数据进行标准化操作？这个问题很难从理论层面解释透彻，即使是BN层的作者给出的解释也未必让所有人信服。与其纠结其缘由，不如通过具体问题来感受数据标准化后的好处。

 考虑Sigmoid激活函数和它的梯度分布，如下图所示，Sigmoid函数在$x\in [-2,2]$区间的导数值在$[0.1,0.25]$区间分布; 当$x>2$或$x<-2$时，Sigmoid函数的导数变得很小，逼近于0，从而容易出现梯度弥散现象。为了避免因为输入较大或者较小而导致Sigmoid函数出现梯度弥散现象，将函数输入x标准化映射到0附近的一段较小区间将变得非常重要，可以从下图看到，通过标准化重映射后，值被映射在0附近，此处的导数值不至于过小，从而不容易出现梯度弥散现象。这时使用标准化手段收益的一个例子。

Sigmoid函数及其导数曲线

 我们再看另一个例子。考虑2个输入节点的线性模型，如图所示:
$$\mathcal{L}=a=x_1{w}_1+x_2{w}_2+b$$
讨论如下两种输入分布下的问题:

• 输入$x_1\in [1,10],x_2\in [1,10]$
• 输入$x_1\in [1,10],x_2\in [100,1000]$

由于模型相对简单，可以绘制出两种$x_1$、$x_2$下，函数的损失等高线图，图（b）是$x_1\in [1,10],x_2\in [100,1000]$时的某条优化轨迹线示意，图（c）是$x_1\in [1,10],x_2\in [1,10]$时的某条优化轨迹线示意，图中的圆环中心即为全局极值点。

数据标准化举例示意图

考虑到:
$$\begin{gathered} \frac{\partial \mathcal{L}}{\partial w_1}=x_1 \\ \frac{\partial \mathcal{L}}{\partial w_2}=x_2 \end{gathered}$$
当$x_1$、$x_2$输入分布相近时，$\frac{\partial \mathcal{L}}{\partial w_1}$、$\frac{\partial \mathcal{L}}{\partial w_2}$偏导数值相当，函数的优化轨迹如图（c）所示; 当$x_1$、$x_2$输入分布差距较大时，比如$x_1\ll x_2$，则:
$$\frac{\partial \mathcal{L}}{\partial w_1}\ll \frac{\partial \mathcal{L}}{\partial w_2}$$
损失函数等势线在$w_2$轴更加陡峭，某条可能的优化轨迹如图（b）所示。对比两条优化轨迹线可以观察到，$x_1$、$x_2$分布相近时图（c）中收敛更加快速，优化轨迹更理想。

 通过上述的两个例子，我们能够经验性归纳出: 网络层输入$x$分布相近，并且分布在较小范围内时（如0附近），更有利于函数的优化。那么如何保证输入$x$分布相近呢？数据标准化可以实现此目的，通过数据标准化操作可以将数据$x$映射到$\hat{x}$:
$$\hat{x}=\frac{x-{\mu }_r}{\sqrt{{\sigma }_r^2+ϵ}}$$
其中${\mu }_r$、${\sigma }_r^2$来自统计的所有数据的均值和方差，$ϵ$是为防止出现除0错误而设置的较小的数字，如$1e-8$。
在基于Batch的训练阶段，如何获取每个网络层所有输入的统计数据${\mu }_r$、${\sigma }_r^2$呢？考虑Batch内部的均值${\mu }_B$和方差${\sigma }_B^2$:
$${\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$$
可以视为近似于${\mu }_r$、${\sigma }_r^2$，其中$m$为Batch样本数。因此，在训练阶段，通过
$${\hat{x}}_{train}=\frac{x_{train}-{\mu }_B}{\sqrt{{\sigma }_B^2+ϵ}}$$
标准化输入，并记录每个Batch的统计数据${\mu }_B$、${\sigma }_B^2$，用于统计真实的全局${\mu }_r$、${\sigma }_r^2$。

 在测试阶段，根据记录的每个Batch的${\mu }_B$、${\sigma }_B^2$估计出所有训练数据的${\mu }_r$、${\sigma }_r^2$，按着
$${\hat{x}}_{test}=\frac{x_{test}-{\mu }_r}{\sqrt{{\sigma }_r^2+ϵ}}$$
将每层的输入标准化。

 上述的标准化运算并没有引入额外的待优化变量，${\mu }_r$、${\sigma }_r^2$和${\mu }_B$、${\sigma }_B^2$均由统计得到，不需要参与梯度更新。实际上为了提高BN层的表达能力，BN层作者引入了“scale and shift”技巧，将$\hat{x}$变量再次映射变换:
$$\tilde{x}=\hat{x}\cdot \gamma +\beta$$
其中$\gamma$参数实现对标准化后的$\hat{x}$再次进行缩放，$\beta$参数实现对标准化后的$\hat{x}$进行平移，不同的是，$\gamma$、$\beta$参数均由反向传播算法自动优化，实现网络层“按需”缩放平移数据的分布的目的。

 下面我们来学习在TensorFlow中实现的BN层的方法。

BatchNorm层实现

1. 向前传播

 我们将BN层的输入记为$x$，输出记为$\hat{x}$。分训练阶段和测试阶段来讨论前向传播过程。

 训练阶段: 首先计算当前Batch的${\mu }_B$、${\sigma }_B^2$，根据
$${\hat{x}}_{train}=\frac{x_{train}-{\mu }_B}{\sqrt{{\sigma }_B^2+ϵ}}\cdot \gamma +\beta$$
计算BN层的输出。

 同时按照
$$\begin{gathered} {\mu }_r\leftarrow \text{momentum}\cdot {\mu }_r+(1-\text{momentum})\cdot {\mu }_B \\ {\sigma }_r^2\leftarrow \text{momentum}\cdot {\sigma }_r^2+(1-\text{momentum})\cdot {\sigma }_B^2 \end{gathered}$$
迭代更新全局训练数据的统计值${\mu }_r$和${\sigma }_r^2$，其中$\text{momentum}$是需要设置一个超参数，用于平衡${\mu }_r$、${\sigma }_r^2$的更新幅度:

当$\text{momentum}=0$时，${\mu }_r$和${\sigma }_r^2$直接被设置为最新一个Batch的${\mu }_B$和${\sigma }_B^2$;

当$\text{momentum}=1$时，${\mu }_r$和${\sigma }_r^2$保持不变，忽略最新一个Batch的${\mu }_B$和${\sigma }_B^2$;

在TensorFlow中，$\text{momentum}$默认设置为0.99。

 测试阶段: BN层根据
$${\tilde{x}}_{test}=\frac{x_{test}-{\mu }_r}{\sqrt{{\sigma }_r^2+ϵ}}\cdot \gamma +\beta$$
计算出${\tilde{x}}_{test}$，其中${\mu }_r$、${\sigma }_r^2$、$\gamma$、$\beta$均来自训练阶段统计或优化的结果，在测试阶段直接使用，并不会更新这些参数。

2. 反向更新

 在训练模式下的反向更新阶段，反向传播算法根据损失$\mathcal{L}$求解梯度$\frac{\partial \mathcal{L}}{\partial \gamma }$和$\frac{\partial \mathcal{L}}{\partial \beta }$，并按着梯度更新法则自动优化$\gamma$、$\beta$参数。

 需要注意的是，对于2D特征图输入$\mathbf{X}:[b,h,w,c]$，BN层并不是计算每个点的${\mu }_B$、${\sigma }_B^2$，而是在通道轴$c$上面统计每个通道上面所有数据的${\mu }_B$、${\sigma }_B^2$，因此${\mu }_B$、${\sigma }_B^2$是每个通道上所有其它维度的均值和方差。以shape为$[100,32,32,3]$为例，在通道轴$c$上面的均值计算如下:

import tensorflow as tf

# 构造输入
x = tf.random.normal([100,32,32,3])
# 将其他维度合并，仅保留通道维度
x = tf.reshape(x, [-1,3])
# 计算其他维度的均值
ub = tf.reduce_mean(x, axis=0)
print(ub)

运行结果如下:

数据有$c$个通道数，则有$c$个均值产生。

 除了在$c$轴上面统计数据${\mu }_B$、${\sigma }_B^2$的方式，我们也很容易将其推广至其它维度计算均值的方式，如图所示:

• Layer Norm: 统计每个样本的所有特征的均值和方差
• Instance Norm: 统计每个样本的每个通道上特征的均值和方差
• Group Norm: 将$c$通道分成若干组，统计每个样本的通道组内的特征均值和方差

 上面提到的Normalization方法均由独立的几篇论文提出，并在某些应用上验证了其相当于或者由于BatchNorm算法的效果。由此可见没深度学习算法研究并非难于上青天，只要多思考、多锻炼算法工程能力，人人都有机会发表创新性成果。

不同标准化方案示意图

3. BN层实现

 在TensorFlow中，通过layers.BatchNormalization()类可以非常方便地实现BN层:

# 创建BN层
layer = layers.BatchNormalization()

与全连接层、卷积层不同，BN层的训练阶段和测试阶段的行为不同，需要通过设置training标志位来区分训练模式还是测试模式。

 以LeNet-5的网络模型为例，在卷积层后添加BN层，代码如下:

network = Sequential([  # 网络容器
    layers.Conv2D(6, kernel_size=3, strides=1),  # 第一个卷积层，6个3×3卷积核
    # 插入BN层
    layers.BatchNormalization(),
    layers.MaxPooling2D(pool_size=2, strides=2),  # 高宽各减半的池化层
    layers.ReLU(),  # 激活函数
    layers.Conv2D(16, kernel_size=2, strides=1),  # 第二个卷积层，16个3×3卷积核
    # 插入BN层
    layers.BatchNormalization(),
    layers.MaxPooling2D(pool_size=2, strides=2),  # 高宽各减半的池化层
    layers.ReLU(),  # 激活函数
    layers.Flatten(),  # 打平层，方便全连接层处理
    layers.Dense(120, activation='relu'),  # 全连接层，120个节点
    # 此处也可以插入BN层
    layers.Dense(84, activation='relu'),  # 全连接层，84个节点
    # 此处也可以插入BN层
    layers.Dense(10),  # 全连接层，10个节点
])

在训练阶段，需要设置网络的参数training=True以区分BN层是训练还是测试模型，代码如下:

with tf.GradientTape() as tape:
    # 插入通道维度
    x = tf.expand_dims(x, axis=3)
    # 向前计算，设置计算模式，[b,784] => [b,10]
    out = network(x, training=True)

在测试阶段，需要设置training=False，避免BN层采用错误的行为，代码如下:

for x, y in test_db:  # 遍历所有训练集样本
    # 插入通道维度，=>[b,28,28,1]
    x = tf.expand_dims(x, axis=3)
    # 向前计算，获得10类别的概率分布，[b,784] => [b,10]
    out = network(x, training=False)

4. 完整代码

加入BN层的LeNet-5完整代码如下:

import os

from Chapter08 import metrics
from Chapter08.metrics import loss_meter

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, Sequential, losses, optimizers, datasets

# 加载MNIST数据集
def preprocess(x, y):
    # 预处理函数
    x = tf.cast(x, dtype=tf.float32) / 255
    y = tf.cast(y, dtype=tf.int32)

    return x, y


# 加载MNIST数据集
(x, y), (x_test, y_test) = keras.datasets.mnist.load_data()
# 创建数据集
batchsz = 128
train_db = tf.data.Dataset.from_tensor_slices((x, y))
train_db = train_db.map(preprocess).shuffle(60000).batch(batchsz).repeat(10)
test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_db = test_db.batch(batchsz)

network = Sequential([  # 网络容器
    layers.Conv2D(6, kernel_size=3, strides=1),  # 第一个卷积层，6个3×3卷积核
    # 插入BN层
    layers.BatchNormalization(),
    layers.MaxPooling2D(pool_size=2, strides=2),  # 高宽各减半的池化层
    layers.ReLU(),  # 激活函数
    layers.Conv2D(16, kernel_size=2, strides=1),  # 第二个卷积层，16个3×3卷积核
    # 插入BN层
    layers.BatchNormalization(),
    layers.MaxPooling2D(pool_size=2, strides=2),  # 高宽各减半的池化层
    layers.ReLU(),  # 激活函数
    layers.Flatten(),  # 打平层，方便全连接层处理
    layers.Dense(120, activation='relu'),  # 全连接层，120个节点
    # 此处也可以插入BN层
    layers.Dense(84, activation='relu'),  # 全连接层，84个节点
    # 此处也可以插入BN层
    layers.Dense(10),  # 全连接层，10个节点
])

# build一次网格模型，给输入x的形状，其中4为随意给的batchsize
network.build(input_shape=(4, 28, 28, 1))
# 统计网络信息
network.summary()

# 创建损失函数的类，在实际计算时直接调用实例即可
criteon = losses.CategoricalCrossentropy(from_logits=True)

optimizer = optimizers.Adam(lr=0.01)

# 训练部分实现如下
# 构建梯度记录环境
# 训练20个epoch


def train_epoch(epoch):
    for step, (x, y) in enumerate(train_db):  # 循环优化
        with tf.GradientTape() as tape:
            # 插入通道维度，=>[b,28,28,1]
            x = tf.expand_dims(x, axis=3)
            # 向前计算，获得10类别的概率分布，[b,784] => [b,10]
            out = network(x, training=True)
            # 真实标签one-hot编码，[b] => [b,10]
            y_onehot = tf.one_hot(y, depth=10)
            # 计算交叉熵损失函数，标量
            loss = criteon(y_onehot, out)
        # 自动计算梯度
        grads = tape.gradient(loss, network.trainable_variables)
        # 自动更新参数
        optimizer.apply_gradients(zip(grads, network.trainable_variables))
        if step % 100 == 0:
            print(step, 'loss:', loss_meter.result().numpy())
            loss_meter.reset_states()

        # 计算准确度
        if step % 100 == 0:
            # 记录预测正确的数量，总样本数量
            correct, total = 0, 0
            for x, y in test_db:  # 遍历所有训练集样本
                # 插入通道维度，=>[b,28,28,1]
                x = tf.expand_dims(x, axis=3)
                # 向前计算，获得10类别的概率分布，[b,784] => [b,10]
                out = network(x, training=False)
                # 真实的流程时先经过softmax，再argmax
                # 但是由于softmax不改变元素的大小相对关系，故省去
                pred = tf.argmax(out,axis=-1)
                y = tf.cast(y, tf.int64)
                # 统计预测样本总数
                correct += float(tf.reduce_sum(tf.cast(tf.equal(pred, y), tf.float32)))
                # 统计预测样本总数
                total += x.shape[0]
            # 计算准确率
            print('test acc:', correct/total)

def train():

    for epoch in range(30):

        train_epoch(epoch)

if __name__ == '__main__':
    train()