基于知识蒸馏的轻量型神经网络设计

郭俊伦，彭书华，李俊杰

(北京信息科技大学 自动化学院，北京 100192)

0 引 言

随着深度神经网络技术的不断发展，深度神经网络在图像分类、目标检测等任务中的识别精度也在不断地提升。2012年，Krizhevsky等[1]提出了深度卷积神经网络AlexNet，证明了神经网络在图像分类任务中的巨大优势。随着研究的不断深入发现，网络的准确率与其深度成正比，但深度加深会引起梯度消失、神经网络退化等问题。2015年，He等[2]提出了残差神经网络，在网络深度达到152层的同时避免了网络加深后可能出现的各种问题。

GPU等硬件设备的不断发展，使得计算机的计算能力以及存储能力都有了很大的提升，但是这也限制了深度神经网络在嵌入式设备上的应用。针对嵌入式设备计算能力差、存储空间小等问题出现了各种网络压缩方法。目前主要的深度神经网络模型压缩算法[3]有5种：轻量型神经网络结构设计、参数剪枝、参数量化、低秩分解、知识蒸馏。

随着神经网络的不断加深，网络的冗余性也随之增加，为了解决这个问题，轻量型神经网络的设计变得尤为重要。轻量型深度神经网络结构设计可使神经网络在深度保持不变的情况下，通过精简网络结构达到减少网络参数、提高准确率以及减少网络冗余性的作用。一般使用分组卷积和轻量化卷积核的方法进行轻量化设计，例如MobileNet[4]、ShuffleNet[5]和ResNext[6]。Hinton等[7]提出了知识蒸馏的概念，在准确率损失较小的情况下，将卷积神经网络从40层减少到16层，大大减小了网络规模，为模型压缩提供了新的研究方向。Zagoruyko等[8]将注意力机制[9]引入知识蒸馏中，将教师网络与学生网络的特征向量图生成的注意力图作为蒸馏介质，提高了学生网络的识别精度。如果将轻量型神经网络与知识蒸馏结合起来，既可以达到缩减网络深度的目的，又可使原网络结构更精简，更加有利于在嵌入式设备上的部署。传统知识蒸馏算法中教师网络与学生网络结构相对独立，对教师网络的深度特征提取能力利用不足，因此本文提出使用改进后的知识蒸馏算法对轻量型神经网络进行知识蒸馏，使设计出的轻量型神经网络在结构尽可能精简的情况下具有更高的精度。

1 理论方法

1.1 轻量型卷积结构

轻量型卷积神经网络的设计是基于残差网络。常规的残差网络中，残差块包含两层卷积层，卷积核为3×3，其基本残差结构如图1(a)所示。为了使网络更加精简，可以使用卷积核为1×1的卷积层代替普通的残差块结构，但这样残差网络的准确率必然会降低，此时可以通过使用知识蒸馏算法来提高网络的准确率。替换方法有两种，其结构分别如图1(b)、(c)所示。

图1 网络的残差块结构

由图1可知，与普通的卷积神经网络相比，残差网络中的卷积神经网络结构是对上一次的计算结果进行修正，而不是直接对上一层的结构进行重新的特征提取，因此可以抑制网络结构产生大的偏差，以解决随着网络加深而产生的梯度弥散以及梯度消失等问题。其原理如式(1)所示：

XL+1=Conv(XL,WL)+XL

(1)

式中：XL、XL+1为残差网络L层、L+1层的运算结果；WL为残差网络结构第L层的参数；Conv为残差块中的普通卷积神经网络的计算函数。

在图1中的普通卷积神经网络部分(两个Conv层)，传统的残差网络一般使用K×K的卷积核进行特征提取，并且往往K>1时神经网络的各个通道才会有信息交互，因此选用K>1的卷积核。设输入维度为N，输出维度为L，那么普通卷积层的计算量为2NKKL。为了尽可能地减少参数数量，一层选用卷积核为1×1的点式卷积层，另一层选用卷积核为3×3的普通卷积层，在卷积层中采用分组卷积的设计方法，分组数为G，此时的网络卷积层计算量为NL+NKKL/G，分组卷积结构如图2所示。

图2 分组卷积结构

轻量型神经网络结构相比于传统残差网络结构参数数量更少，部署难度更低，但是在一定程度上却损失了准确率，因此需要结合改进的知识蒸馏算法以提高其准确率。

1.2 改进的知识蒸馏算法

在Softmax中引入温度系数T，并且将训练标签由硬标签转变为教师网络Softmax输出的软标签，此时知识蒸馏的Softmax公式为

(2)

式中：qi为第i类的概率；zi为神经网络结果第i维的输出结果；exp为指数函数。

学生模型的损失函数为

Lkd=αT2Cr(qs,qt)+(1-α)×Cr(qs,y)

(3)

式中：α为蒸馏系数，即软标签在训练学生模型时所占的比重；T为温度系数，可使软标签更加平滑；qs、qt分别为引入温度系数T后学生网络和教师网络的Softmax公式；Cr为交叉熵函数；qs和y为传统训练方法中学生网络的Softmax公式以及硬标签。

传统的知识蒸馏架构中学生网络与教师网络相对来说较为独立，没有充分利用到教师网络的特征提取能力。针对这个缺点，提出并使用了一种结合图片风格迁移实验[10]中的风格迁移算法以及传统知识蒸馏算法的改进版知识蒸馏算法。改进后的算法不仅使用了传统的知识蒸馏公式，并且利用了教师网络强大的特征提取功能，结合了风格迁移实验中的特征重建以及图片风格迁移算法作为特征评估器，用来矫正学生网络训练时产生的误差。

2 神经网络设计

传统的知识蒸馏算法是教师网络和学生网络都选用结构相似，主要是模型宽度相似的网络进行知识蒸馏，比如基于注意力机制的知识蒸馏方法中使用的是以宽度为基准的宽度残差神经网络。在这种网络中，学生网络和教师网络进行知识蒸馏的中间层具有相同的宽度，方便进行两个网络中间特征层的特征对比。

在设计基于知识蒸馏的轻量型神经网络训练结构时，考虑到分组卷积最后的堆叠排列有一定的人为设计因素的影响，所以针对这种情况选取的训练结构主要有以下两种情况：1)教师网络与学生网络结构考虑到人为因素的影响，采用基于风格迁移的知识蒸馏架构；2)由于神经网络具有自适应性，所以不考虑人为因素的影响，仍然采用基于特征重建的知识蒸馏算法[11]。

2.1 基于风格迁移的知识蒸馏算法

针对第一种情况，教师网络与学生网络为轻量型神经网络，也就是使用分组卷积结构，卷积层的输出特征向量的分布是直接堆叠而成的，有一定的人为设计因素。由于教师网络和学生网络本身的特征提取能力有较大的差异，这使得两种网络有不同的特征排列方式，所以选用结合风格迁移实验中风格迁移的知识蒸馏算法，其结构如图3所示。

图3 基于风格迁移的知识蒸馏算法

在风格迁移实验中主要是图片风格迁移算法(即图片背景迁移)和特征重建算法(即图片内容迁移)两种算法。图片风格越相似则深度特征的Gram矩阵就越相似，图片内容越相似则网络的特征越相似。特征重建算法和风格迁移算法分别如式(4)、式(5)所示。

(4)

(5)

基于风格迁移的知识蒸馏算法与传统的知识蒸馏算法结合后可以得到最终的损失函数：

(6)

式中：λ为总体损失函数中风格迁移算法部分所占的权重。

2.2 基于特征重建的知识蒸馏算法

针对第二种情况，不考虑人为设计因素的影响，只考虑神经网络的自适应性，采用基于特征重建的知识蒸馏算法，其结构如图4所示。

图4 基于特征重建的知识蒸馏算法

基于特征重建的知识蒸馏算法结合传统知识蒸馏算法后的损失函数为

(7)

式中：β为总体损失函数中风格迁移算法部分所占的权重。

3 实验结果及分析

3.1 轻量型神经网络实验结果

轻量型神经网络设计基于宽度残差网络结构(wide residual net，WRN)，表1～4中的WRN-16-2表示网络深度为16，网络宽度为基准宽度的2倍。将原网络的卷积层平均分为3组，每组宽度由浅到深依次为16、32、64。为了对比普通分组卷积与点式卷积先后顺序对网络性能的影响，在不同复杂度的数据集Cifar10与Cifar100上进行实验。

由Crowley等[12]的实验可以发现轻量型神经网络的准确率往往与分组数成反比，本实验中分组数选用中间值4。为了验证点式卷积在第一层或第二层(即残差块的不同设计方式)对整体网络结构的影响，设计了4种不同类型的网络结构进行对比实验。

表中PG4表示残差块结构中第一层为点式卷积(pointwise convolution，简称P)，第二层为分组卷积(grouped convolution，简称G)，使用大小为3×3的卷积核；G4P与PG4的一、二层顺序相反；P4G4表示残差块结构中第一层为分组点式卷积，使用大小为3×3的卷积核，第二层为点式分组卷积；G4P4与P4G4的一、二层顺序相反。参数数量为WRN-16-2网络的参数数量。

在Cifar10数据集的实验结果如表1所示。由表1可知，在Cifar10数据集上4种结构中G4P(即图1(c)中所示的网络结构)稳定性最好，同时具有较高的准确率。

表1 四种类型轻量型网络在Cifar10上的错误率及参数数量

为了验证不同网络结构在复杂任务中的分类能力，选用Cifar100继续验证网络的适用性。其实验结果如表2所示。

表2 四种类型轻量型网络在Cifar100上的错误率及参数数量

对表2分析可知，点式卷积本身就会削弱卷积神经网络通道间的信息交互，若再进行分组卷积则势必会降低网络的准确率，并且G4P这种网络结构还具有参数数量相对较低的优势，因此选择G4P作为最终的设计。

3.2 结合知识蒸馏算法后的实验结果

通过在Cifar10与Cifar100中的对比实验发现，先分组后点式卷积的结构G4P具有一定的优势，但是由于这种轻量型结构会降低模型的精度，故可采用知识蒸馏算法对网络的准确性进行提升。以下实验中学生网络结构为WRN-16-2，教师网络结构为WRN-40-2，传统知识蒸馏算法(式(3))选用参数T=4，α=0.9。基于风格迁移的知识蒸馏(ST+KD)算法(式(6))中，使用参数λ=0.001；基于特征重建的知识蒸馏(FR+KD)算法(式(7))中，使用参数β=0.01。

在Cifar10上的实验结果如表3所示。由表3可知，使用不同的知识蒸馏架构均可以提升网络的准确率，而两种不同的架构准确率仅仅相差0.15%，由此可知由于网络具有自适应性，可以忽略人为设计因素产生的影响。

表3 轻量型神经网络在Cifar10上的错误率

在Cifar100上进行实验的结果如表4所示。

表4 轻量型神经网络在Cifar100上的错误率

实验结果表明，知识蒸馏算法可以提升轻量型神经网络的准确性，并且两种架构的准确率仅仅相差0.19%，所以可以忽略人为设计因素带来的影响。

3.3 轻量型神经网络的优势分析

经过知识蒸馏算法设计的轻量型神经网络WRN-16-2在Cifar10的准确率为92.38%，网络参数数量为1.317×105；在Cifar100的准确率为68.6%，网络参数数量为1.433×105。而在传统的知识蒸馏算法实验中，使用普通残差网络结构的WRN-16-1在Cifar10的准确率为92.03%，网络参数数量为1.759×105；在Cifar100的准确率为67.26%，网络参数数量为1.818×105。由此可见采用轻量型神经网络设计不仅提升了准确率，而且使网络参数降低了20%以上。

4 结束语

采用分组卷积以及点式卷积结合的残差块结构构成轻量型神经网络基本结构，经过训练，在参数相同的情况下模型准确率要高于使用普通卷积结构的残差网络。使用基于特征重建以及基于风格迁移的两种架构均可提升网络的准确率，且经过对比实验发现由于神经网络的自适应性可以忽略分组卷积算法等一定的人为设计因素。由于蒸馏架构以及轻量型神经网络均有人工设计的痕迹，这在一定程度上限制了神经网络的自适应性，虽然其影响较小，但是不能完全忽略，并且知识蒸馏算法选用不同的蒸馏架构时其效果也存在不确定性。所以该方法具有一定的局限性，需进一步完善。