基于知识蒸馏的特征精炼相互学习方法

文 强，郭 涛，王 涛，李贵洋，邹俊颖

(四川师范大学 计算机科学学院，四川 成都 610101)

0 引 言

神经网络在解决很多复杂问题上的能力有了指数级的进步，已经应用于计算机视觉的各个领域[1-3]。由于训练模型包含大量参数，部署运行需要占用较大存储空间且执行速度缓慢，限制了其在移动设备或嵌入式设备等低内存或有快速执行要求的应用程序和平台中的部署与运行[4，5]。知识蒸馏是一种有效加速模型速度和压缩模型大小的方法，在模型大小和准确率上能取得很好的平衡，逐渐成为研究热点。

知识蒸馏是将知识从预训练的教师网络转移到未经训练的学生网络的模型压缩方法，在部署阶段，较小的学生网络取代大型教师网络。传统的知识蒸馏通常采用单向知识转移和两阶段训练程序，即提前训练好教师网络，学生网络再从教师网络单向获取知识。然而，复杂容量的教师模型需要较长训练时间且占用大量的存储。为了减少训练大容量高性能的复杂网络和在没有预训练的教师网络等情况下，提高小型轻量型网络的准确率以及降低模型大小，自我知识蒸馏和相互蒸馏等方法得到了快速发展。

本文提出一种基于知识蒸馏的特征精炼相互学习方法(feature refining mutual learning，FRML)。该方法在没有教师网络的情况下，给相互学习的学生网络配备辅助网络，通过对处理多尺度特征网络的调整达到对知识提炼目的，生成精炼特征图。辅助网络将精炼的知识传递给其辅助的学生网络，而学生网络之间通过协作训练，学习互相精炼后的知识，优化自身网络，提高网络的泛化能力以及其对视觉任务的适用性。

1 相关工作

1.1 知识蒸馏

知识蒸馏的主要思想源于Buciluǎ等提出的模型压缩[6]，即将信息从大模型或模型集合转移到训练小模型，使得小模型近似大模型，从而不显著降低准确率。知识蒸馏是一种类似于迁移学习的模型压缩方法，但知识蒸馏不同于迁移学习。知识蒸馏的目标是提供较小的网络模型来解决与较大网络模型相同的任务；而迁移学习的目标是减少网络模型的训练时间，以解决类似于其它网络模型所解决的任务。知识蒸馏通过改变正在训练学生网络的损失函数来实现其目标，以解释预训练教师网络的隐藏层输出；而迁移学习通过将预训练网络模型的参数迁移到新模型来实现其目标。传统知识蒸馏传递知识的方法为，给定一个logitsz向量作为教师网络的最后一个全连接层的输出，zi是第i类的logit，采用softmax函数估计输入属于第i类的概率pi， 引入温度因子T控制每个软标签的重要性，如式(1)

(1)

T越高在各类别上产生越弱的概率分布。通过以上方法，教师网络获得的软标签预测包含暗知识，并且可以用作监督，以将知识从教师网络转移到学生网络。教师网络提供不同的软标签会给同一个学生网络带来不同的表现。

1.1.1 自我知识蒸馏

自我知识蒸馏不同于传统知识蒸馏，它不需要提前准备大型复杂的教师网络，即学生网络利用自身知识来提高训练自身网络的有效性。Zhang等[7]提出在只有学生网络的情况下，将深层网络的部分知识蒸馏到网络较浅部分，用中间隐藏层特征对输出进行分类。Hou等[8]提出用于通道检测的自注意蒸馏方法，该网络利用自身各层的注意力图作为下层蒸馏目标，以减少模型参数量。Zhu等[9]为了降低计算成本，提出添加额外分支，使模型参数和中间隐藏层的估计特征多样化，其中每个分支表示一个学生网络，最后集成所有的分支网络进行输出。上述自我知识蒸馏方法基本上没有对网络进行增强，导致缺乏更精炼知识。Ji等[10]提出在理想的情况下，可以通过从复杂集成模型中提取精炼知识以改进特征图蒸馏。

1.1.2 相互蒸馏

相互蒸馏在没有复杂预训练的教师网络情况下，教师网络和学生网络同时学习，教师网络会随着学生网络的更新而更新。Zhang等[11]提出深度相互学习，在训练过程中，任何一个网络均可作为学生，其它网络作为老师，相互教学以提高模型准确率。Chen等[12]提出将辅助同伴(auxi-liary peers)和组长(group leader)引入深度相互学习中，以形成一组多样化的同伴网络。Kim等[13]提出采用特征融合模块构建教师网络，将多个并行神经网络训练成子网络，利用融合模块将子网络知识结合起来，传递给教师网络，教师网络再将其知识传递给每个子网，以在线知识提取的方式相互教学。但是这种相互教学的方式，在训练过程中会丢失局部信息，阻碍了其对各种视觉任务适用性，这种知识提炼方法，单个教师网络或者学生网络均接收不到精炼的特征图。

1.2 特征网络

特征网络又叫特征金字塔网络[14](feature pyramid networks，FPN)，是一种以任意尺寸的单尺度图像为输入，以全卷积方式在多个层次上输出按比例大小的特征图的特征提取器。自从FPN被提出来以后，均只有自上向下的单向融合。Liu等[15]提出路径聚合网络(path aggregation network，PANet)，为FPN引入额外自下而上的融合路径，以实现网络中层和层之间的短路径连接，从而实现二次融合。Tan等[16]提出一种加权双向特征金字塔网络(weighted bi-directional feature pyramid network，BiFPN)，该网络能简单快速地进行多尺度特征融合。BiFPN采用与PANet相同的自上向下和自下向上网络结构，同时增加部分跳跃连接，并将该网络结构融合构造为一个模块，使其可以重复堆叠，增强特征融合。同时，BiFPN为每个输入特征增加额外权重，允许网络学习每个输入特征的重要性，从而获得精炼的特征。

2 FRML模型

2.1 基本定义

定义1 数据集：训练图像数据集D={(x1，y1)，(x2，y2)，…，(xN，yN)}， 其中N为样本大小，yi∈{1，2，…，M}，M表示类别大小。

定义2 学生网络：Θ1，Θ2为学生网络，并且Θ1=Θ2，Θ1有n个块。设Sij为样本i输入Θ1的第j个块的特征映射，cj为Θ1的第j个块的通道维度。

定义3 辅助网络：Θ′1、Θ′2分别为Θ1、Θ2的辅助网络，并且Θ′1=Θ′2，Θ′1有n个块。设Tij为输入特征映射，为了简化符号，在本文的其余部分省略了索引i。学生网络与辅助网络的特征映射块数量相同。

定义4 logit与预测值：logitz1和z2分别为Θ1和Θ2的完全连接层的输出；z′1和z′2分别为Θ′1和Θ′2完全连接层的输出；p1和p2分别为Θ1和Θ2输出预测值；p′1和p′2分别为Θ′1和Θ′2输出预测值。

2.2 模型结构

FRML模型由两个对称的学生网络Θ1和Θ2以及其辅助网络Θ′1和Θ′2构成。首先图像分别输入该模型的Θ1和Θ2， 提取到n个特征图块Sj， 并输出特征图z1，z2； 其次将每个特征图块Sj分别输入到Θ′1和Θ′2， 得到精炼的特征图块Tj， 并输出特征图z′1和z′2； 根据每个输出特征图z1，z′1，z2，z′2， 计算得到每个输出预测值p1，p′1，p2，p′2。 最后Θ1和Θ2根据自身的特征图块Sj与Θ′1和Θ′2的特征图块Tj、 输出预测值、真实标签，做出相应计算，得到每个部分的损失，并计算总损失以及更新权重。两个学生网络相互训练，相互学习到精炼的特征，从而提高模型的性能。FRML结构如图1所示。

图1 FRML结构

2.2.1 学生网络

FRML中两个学生网络Θ1和Θ2采用相同的ResNet[17]网络，该网络包含4层block，使用每个阶段的最后一个残差结构的特征图作为输出，表示为 {S1，S2，S3，S4}， 通过完全连接层输出特征图z1和z2， 完成每个阶段的特征提取，如图1所示。

Θ1和Θ2同时训练，相互传递知识，进行交流。样本xi输入Θ1之后所属类别m的概率通过式(2)计算

(2)

对于多分类任务，采用交叉熵函数计算预测值与真实标签的误差，如式(3)、式(4)

(3)

指示函数I定义为

(4)

为提高Θ1在测试实例上的泛化性能，引入与Θ1相互学习的Θ2， 以Θ2后验概率p2的形式提供训练经验。为量化Θ1和Θ2的预测值p1和p2， 使用Kullback Leibler(KL)散度。从p1到p2的KL散度计算如式(5)

(5)

至此，Θ1和Θ2相互学习损失LΘ1和LΘ2分别为式(6)、式(7)

LΘ1=LCE1+DKL(p2‖p1)

(6)

LΘ2=LCE2+DKL(p1‖p2)

(7)

两个学生网络通过相互学习的方式，每个学生网络均能学习训练实例的真实标签，即学生网络受监督损失LCE指导以及匹配其一起相互学习的概率估计，即KL散度损失。

2.2.2 辅助网络

为帮助Θ1提供精炼的特征映射和软标签，给相互学习的Θ1添加Θ′1。 假定Θ1有n个块，辅助网络的输入为学生网络的特征映射S1，…，Sn。 辅助网络采用来自PANet[15]自下而上的路径和FPN[21]自上而下的路径，以适应分类任务。辅助网络结构如图2所示。

图2 辅助网络结构

在进行自上而下的路径之前，利用式(8)的横向卷积层将特征映射输入进辅助网络

Lj=Conv(Sj；dj)

(8)

Conv为输出维度为dj的卷积运算。dj取决于Θ1输出特征映射的通道维度cj。

Sj经过横向卷积得到Lj后，通过自上而下路径，融合不同层次特征，提取低层信息Pj， 再通过自下而上路径，进行二次特征融合，得到Tj， 如式(9)和式(10)所示

(9)

(10)

通过图2可以看出，首先，如果节点仅包含输入，并没有进行特征融合，那么它对于特征网络的贡献是极小的。因此，将P4和P1移除，从而提高模型计算效率。其次，针对同一层级，添加一条从起始输入到输出的连接，即L2到T2和L3到T3， 在不增加任何参数的同时，融合更多特征。最后，添加L4到P3和P2到T1的连接，创建完整的自上而下路径。其中，{P2，P3} 和 {T1，T2，T3，T4} 通过3×3 Conv得到。对于参数wP和wT， 使用快速归一化融合技术[16]，引入可学习的权重，用于学习不同输入特征的重要性。Resize运算符分别代表自上而下路径上的上采样和自下而上路径上的下采样；对上采样使用双向线性插值，即融合上一层2×UpSample的特征，对下采样使用最大池化，即融合下一层3×3/2 Maxpool的特征。最后在自下而上路径的顶部增加全连接层输出z′1， 增加softmax层来输出p′1预测类别。

(11)

与其它自我知识蒸馏方法类似，FRML通过软标签进行蒸馏，提供软标签蒸馏损失，如式(12)所示

(12)

其中，T是温度系数。此外，Θ′1使用交叉熵损失函数学习真实标签，通过式(3)计算得出L′CE1。

2.2.3 损失函数

(13)

(14)

其中，α和β为超参数，本文选择α∈[1，2，3] 和β∈[100，200]。 对于两个学生网络及其辅助网络，通过反向传播算法进行优化。FRML模型通过交叉熵损失、KL散度损失、自特征蒸馏损失和软标签蒸馏损失更新参数。

2.3 算法流程

精炼特征相互学习方法分为两个阶段：首先，两个学生网络通过辅助网络分别获取到精炼的知识，然后，两个学生网络相互学习精炼的特征，两个学生网络互相微调，以提高自身的泛化能力。FRML模型的算法的伪代码如算法1所示。

算法1：特征精炼相互学习算法

输入：学生网络Θ1和Θ2； 辅助网络Θ′1和Θ′2； 训练数据集；学习率、小批量尺寸以及α，β

输出：更高分类准确率的学生网络

(1)随机初始化Θ1和Θ2中所有网络层的参数

(2)for 每个小批量

2.1 输入样本xi, 真实标签yi;

2.2 输入Θ1, 将Sj输入到Θ′1, 得到Tj并得到特征图z1,z′1;

2.3 输入Θ2, 将Sj输入到Θ′2, 得到Tj并得到特征图z2,z′2;

2.4 特征图通过式(2)分别得到预测值p1、p′1和p2、p′2;

2.5 通过式(3)计算交叉熵损失LCE1, L′CE1和LCE2, L′CE2;

2.10 计算Θ1的总损失

2.11 使用反向传播算法更新Θ1的权重;

2.12 计算Θ2的总损失

2.13 使用反向传播算法更新Θ2的权重;

(3)end for

3 实验结果及分析

3.1 实验数据集

本文使用了6种数据集，其中CIFAR100和TINY由小比例图像组成，这里将TINY图像的尺寸调整为与CIFAR100图像尺寸相同，即32×32大小。CUB200、MIT67、Stanford40和Dogs均为细粒度视觉分类(fine-grained visual categorization，FGVC)任务的图像数据集。实验使用数据集概况见表1。

表1 本文所用数据集概况

3.2 实验设置

FRML模型在数据集CIFAR100和TINY上采用相同的ResNet18[17]和WRN-16-2[20]作为学生网络。为了使ResNet18适应小型数据集，将ResNet18的第一个卷积层修改为3×3的内核大小，采用单个步长和单个填充，并删除最大池化操作。在FGVC任务数据集上采用标准的ResNet18作为学生网络。

实验基于深度学习框架Pytorch 1.8.0、Python 3.8.8 实现，并在NVIDIA GeForce RTX 3090 GPU上完成。

对于所有训练样本，使用随机裁剪和翻转来增加数据的多样性。在训练过程中，采用随机梯度下降(SGD)方法，初始的学习率为0.099，动量参数为0.9，优化器中权重衰减参数为0.0001；温度系数T为4；通道宽度参数w为2；超参数[10]的设置见表2。

表2 超参数设置

3.3 分类性能对比实验

对于图像分类任务，采用测试集Top-1准确率(Accuracy)作为评价指标。本实验完成了FRML模型与Baseline、DML[11]、FRSKD[10]这3种分类方法在CIFAR100、TINY和FGVC任务数据集上的准确率比较。实验结果见表3和表4。

表3 FRML在CIFAR100和TINY数据集上的分类准确率

表4 FRML在FGVC任务数据集上的分类准确率

从表3可知，FRSKD和DML均提高了标准分类器网络的性能，但FRML表现出更好的性能，与基线FRSKD相比，平均准确率提高了1.49%，与基线DML相比，平均准确率提高了2.30%。

表4显示了在FGVC任务数据集上的分类准确率。在4种数据集上的实验结果可知，FRML表现出比其它方法更好的性能。与基线FRSKD相比，平均准确率提高了1.85%，与基线DML相比，平均准确率提高了8.08%。通过实验验证，FRML提高了相互学习在细粒度视觉任务的准确率。

3.4 知识蒸馏分类实验

在知识蒸馏分类实验中，将FRML与FitNet[21]、ATT[18]、Overhaul[19]等知识蒸馏方法做对比实验。教师网络设置为预训练的ResNet34，学生网络设置为未经训练的ResNet18。其中，分类准确率的最大值用粗体标识，实验结果见表5。

表5 知识蒸馏上分类准确率

从表5中可知，来自强大的教师网络的传统知识蒸馏方法虽然提高了学生的表现。但FRML在所有数据集上优于使用教师网络的知识蒸馏方法，这意味着在相互学习的过程中，通过精炼特征传递给学生网络，比预训练的教师网络更好。

3.5 扩展不同学生网络的实验

为了进一步提高模型的适用性，不限制模型架构，本实验将两个学生网络扩展成不同结构的网络。其中，在数据集CIFAR100和TINY上分别采用WRN-16-2和ResNet18作为学生网络Θ1和Θ2， 在FGVC任务数据集上分别采用ResNet18和ResNet32作为学生网络Θ1和Θ2。 FRML模型在不同网络架构以及不同数据集上分类准确率见表6。

表6 不同学生网络分类准确率

从表6可知，相比较独立训练的网络和DML，基于不同网络的组合，FRML均提高了每个学生网络的性能。与DML相比，在数据集CIFAR100和TINY上，WRN-16-2网络的平均准确率提高了1.65%，ResNet18网络的平均准确率提高了2.00%。在FGVC任务数据集上采用ResNet18网络的平均准确率提高了3.84%，采用ResNet32网络的平均准确率提高了4.03%。由此可知，不同网络架构或不同大小的网络，通过特征精炼相互学习后，性能均得到提高。实验结果表明，在不同网络架构上，本方法同样能够改进相互学习的不足，提高不同网络的泛化性能。

3.6 模型收敛性实验

本实验为FRML模型在FGVC任务数据集上训练次数以及测试分类准确率上的收敛实验，结果如图3所示。图3(a)～图3(d)分别为模型在CUB200、MIT67、Dogs、Stanford40数据集上的收敛效果。

图3 FRML在FGVC任务数据集上的收敛曲线

从图3(a)～图3(d)可以看出，FRML模型在FGVC任务数据集上均具有良好的收敛性。当训练次数在100次之前，准确率的提升速度较快；在训练到100次左右时有较大提升；在之后的训练次数中，分类准确率小幅度振动，模型开始整体进入收敛状态；训练到150次左右时，可以看出准确率有个小幅度的上升；在150次之后，分类准确率波动趋于稳定，模型基本达到收敛。

此外，从图3(a)～图3(d)可以看出，对比FRML和DML，FRML在FGVC任务数据集上的分类准确率明显高于DML。实验结果表明本方法解决了相互学习上的不足，提高了学生网络在相互学习的泛化性能，以及其对视觉任务的适用性。

4 结束语

本文在没有预训练的教师网络的情况下，提出一种基于知识蒸馏的特征精炼相互学习方法，解决两个学生网络在相互学习时出现图像局部信息丢失，而造成单个学生网络接收精炼特征图不完整的问题。通过对每个学生网络配备特征提取辅助网络进行特征融合，利用辅助网络输出精炼特征映射和软标签对学生网络进行训练，并保留图像特征的局部信息，为学生网络之间传递精华知识，以达到最优效果。通过改进相互学习策略，提高了视觉任务的适用性。后续的研究是将学生网络的数量扩展到3个及以上以进一步提高其性能。