CN115063663A - 一种基于知识蒸馏的目标检测压缩方法 - Google Patents

Abstract

本发明公开了一种基于知识蒸馏的目标检测压缩方法，包括：S1、对所有图像进行放缩操作；S2、对图像进行预处理，然后将其进行拼接；S3、对教师网络进行训练，训练过程中，针对教师网络中的骨干网络蒸馏损失、neck部分蒸馏损失以及检测头蒸馏损失进行加权；S4、训练后，冻结教师网络参数，然后再次将图像送入教师网络中获得输出；S5、将输出结果连同标签一起送入学生网络中，得到蒸馏损失和目标检测损失后进行反向传播，最小化损失函数得到最终蒸馏后的学生网络，即得到压缩后的目标检测模型。本发明可以有效提升检测的精度，并降低模型的参数量，从而更好地应用于算力更小的边缘设备，为社会带去更多科技的福利。

Description

一种基于知识蒸馏的目标检测压缩方法

技术领域

本发明涉及工业物联网技术领域，特别涉及一种基于知识蒸馏的目标检测压缩方法。

背景技术

近年来，随着ImageNet(用于视觉对象识别软件研究的大型可视化数据库)的出现以及GPU设备算力的增强，计算机视觉领域得到了极大的发展，尤其是以深度学习为代表的一系列卷积神经网络的出现更是极大的推动了该领域的进步。随着深度学习技术的不断发展，网络结构也朝着更深更宽的方向发展，这也间接导致了极大的计算和显存开销。

基于此，有人提出了模型压缩方法知识蒸馏。它是一种特殊的知识迁移方法，其主要思想是通过将较大模型的输出交由较小模型进行学习，通过模仿输出结果达到类似于知识传递的效果，因此这种方法也被称为教师学生网络。目前，知识蒸馏已经是一种极为有效的网络训练手段，它能够有效地提高模型精度，同时增强网络的泛化性能，避免计算和显存开销的增加。

然而，尽管目标检测有着广泛的应用场景，但在一些资源受限或边缘设备场景依然没法大规模的应用，例如个人手机、小型嵌入式设备等。因此，有必要针对现有的模型压缩方法进行改进，以便能够更好地适用于算力更小的边缘设备。

发明内容

本发明的目的在于提供一种基于知识蒸馏的目标检测压缩方法，可以有效提升检测的精度，并降低模型的参数量，从而更好地应用于算力更小的边缘设备，为社会带去更多科技的福利。

为实现上述目的，本发明采用的技术方案如下：

一种基于知识蒸馏的目标检测压缩方法，包括以下步骤：

步骤S1、对所有图像进行放缩操作，并确保图像缩小后比例与原始比例一致；

步骤S2、对图像进行预处理，然后将其进行拼接；

步骤S3、利用数据集中的图像及标签，结合随机梯度下降方法对教师网络进行训练，训练过程中，针对教师网络中的骨干网络蒸馏损失、neck部分蒸馏损失以及检测头蒸馏损失进行加权，其中，骨干网络蒸馏损失中包含有全局特征权重、局部特征权重以及带权特征权重三个部分；

步骤S4、训练后，冻结教师网络参数，然后再次将图像送入教师网络中获得输出；

步骤S5、将输出结果连同标签一起送入学生网络中，得到蒸馏损失和目标检测损失后进行反向传播，最小化损失函数得到最终蒸馏后的学生网络，即得到压缩后的目标检测模型。

所述步骤S1具体为：将所有图像进行放缩操作，使其最长边不超过1333像素，最短边不低于800像素，在放缩的同时避免图片长宽比失真，确保图像缩小后比例与原始比例一致。

具体地，所述步骤S2包括以下步骤：

步骤S201、将图像随机翻转，翻转的概率控制在0.5；

步骤S202、对图像像素值进行正则化处理；

步骤S203、将放缩后的图像进行填充，使其横纵坐标的像素总数能被32整除；

步骤S204、将同批次的图像整体处理后，扩充第一维通道，在第一维通道上进行拼接。

进一步地，所述步骤S202中采用的正则化的计算公式为：

像素输出＝(像素值-均值)/方差

其中，三通道的均值和方差分别是[123.675,116.28,103.53]及[58.395,57.12,57.375]，该数值是由COCO训练集整体计算所得。

再进一步地，所述步骤S3中，训练的参数设置如下：

学习率为0.01，动量设为0.9，权重衰减设为0.0001，训练过程中采用线性暖启动策略，迭代数量设为500，在周期为8和11时对学习率进行减小，总共训练12个epoch；

检测头的分类损失采用focal loss，其中β参数为2，α参数为0.25；回归损失为一范数损失。

再进一步地，所述步骤S3中，随机梯度下降方法的损失函数为：

式中，loss(·)表示损失函数，x_i表示标签的真实分布，y_i表示模型的预测分布，n表示标签的数量，在分类任务中采用交叉熵表示，即：

回归任务中采用均方误差表示，即：

再进一步地，所述步骤S3中，蒸馏过程的激活函数为：

式中，z_i为第i个节点的输出值，z_j为第j个节点的输出值，T为蒸馏温度。

再进一步地，骨干网络蒸馏损失公式为：

式中，α＝0.1，β＝0.05，γ＝1000，p(x_i)表示真实概率分布，q(x_i)表示模型预测的概率分布；

蒸馏的总体损失公式为：

loss_dis＝λ₁loss_{backbone_dis}+λ₂loss_{neck_dis}+λ₃loss_{head_dis}

式中，loss_{backbone_dis}、loss_{neck_dis}、loss_{head_dis}分别代表骨干网络蒸馏损失、neck部分蒸馏损失以及检测头蒸馏损失，λ₁、λ₂、λ₃用于平衡各个蒸馏损失之间的权重，λ₁＝1，λ₂＝0.05，λ₃＝0.3。

再进一步地，所述步骤S3中，网络总体损失的计算公式为：

loss＝loss_det+η·loss_dis

式中，

是交叉熵损失函数，

是smooth L1函数，η表示蒸馏损失权重，λ是归一化权重，t_i是一个向量，表示在RPN训练阶段，anchor相对于gt预测的偏移量，

是与t_i维度相同的向量，表示在RPN训练阶段，anchor相对于gt实际的偏移量，p_i为anchor预测为目标的概率，

有物体时为1，没有物体时为0，N_cls是总的anchor数量，N_reg是featuremap的size。

再进一步地，蒸馏损失权重调整策略的计算公式为：

式中，cur_iter表示当前训练的迭代次数；max_iter为超参数，表示蒸馏损失调整的最大迭代次数。

本发明的主要设计原理在于：使用Pytorch深度学习框架，引入注意力机制并采取新颖的蒸馏方式对不同种类的目标检测算法分别进行蒸馏，这种蒸馏方式采用了多尺度特征融合表征，可以解决蒸馏过程中特征不突出、无效蒸馏的问题，“多尺度特征融合蒸馏”整体框架如图1所示，“卷积块注意模块”注意力模块如图2所示。

具体地，本发明采用的蒸馏方式主要是通过教师网络的中间层特征生成通道及空间特征响应并得到对应特征掩码，学习教师网络中在输出具有较大响应的显著特征。同时，将知识回顾的方式应用于该算法中，通过不同层次特征之间的融合交互，提高网络的蒸馏效率。

并且，为了提高网络整体的蒸馏效率，在neck及检测头部分，该蒸馏方式采用了掩码引导的蒸馏方法，通过教师网络检测头的分类输出得到特征的重要性权重，利用该权重作为掩码的重要性依据分别对neck及检测头的输出进行蒸馏。

与现有技术相比，本发明具有以下有益效果：

(1)针对两阶段目标检测，本发明以ResNet50作为骨干网络的FasterRCNN网络，在经过ResNet101网络的蒸馏后，网络mAP整体相较于未蒸馏模块可以获得有效的提高(经验证，可提高1.3个百分点)，尤其是在小目标的指标mAPs上，超过了教师网络0.3个百分点。同时，经过蒸馏后的学生网络精确度与教师网络相当，仅相差0.7个百分点，但FLOPs及Parameters指标却只有ResNet50作为骨干网络时的指标，大幅减少了模型的参数量。

(2)针对单阶段目标检测，本发明经过蒸馏后以ResNet50为骨干网络的RetinaNet算法相较于未蒸馏前的算法同样有显著的提升(经验证，各项指标平均上升了1.5个百分点)，同教师网络相比也仅仅相差了0.3个百分点，并在精度少量降低的情况下，原本只能跑7.07帧的算法跑到了8.78帧，同时模型大小从56.74M降低至37.74M，不仅有效节约了模型计算需要占用的空间，而且计算效率也得到了较好的提高。

(3)针对无锚框算法FCOS，本发明经过蒸馏后的学生网络，经验证，其精确度超过了教师网络0.2个百分点，且mAP指标达到了39.3。在其他小指标中，mAPs和mAP75指标也都超过了教师网络，同时在相同骨干网络下，拥有最少的浮点计算次数和参数量，并且在FPS指标上首次超过了10。

综上，本发明通过使用Pytorch深度学习框架，引入注意力机制并采取新颖的蒸馏方式对不同种类的目标检测算法分别进行蒸馏，大幅提高了检测模型最终的准确性，并且降低了模型的参数量及提高了模型的计算效率，因此也更加适用于诸如个人手机、小型嵌入式设备等算力更小的边缘设备。

附图说明

图1为本发明采用的多尺度特征融合蒸馏整体框架示意图。

图2为本发明采用的卷积块注意模块注意力模块示意图。

图3为本发明-实施例的流程示意图。

具体实施方式

下面结合附图说明和实施例对本发明作进一步说明，本发明的实施包含但不限于以下实施例。

实施例

本实施例提供了一种基于知识蒸馏的目标检测压缩方法，旨在能够更好地适用于算力更小的边缘设备(如个人手机、小型嵌入式设备)。下面对本实施例的实施流程进行详细介绍。

如图3所示，首先，将所有图像进行放缩操作，使其最长边不超过1333像素，最短边不低于800像素，在放缩的同时避免图片长宽比失真，确保图像缩小后比例与原始比例一致。

放缩图像后，对图像进行预处理，在本实施例中，预处理的流程如下：

(1)将图像随机翻转，翻转的概率控制在0.5；

(2)对图像像素值进行正则化处理，正则化的计算公式为：

像素输出＝(像素值-均值)/方差

其中，三通道的均值和方差分别是[123.675,116.28,103.53]及[58.395,57.12,57.375]，该数值是由COCO训练集整体计算所得；

(3)将放缩后的图像进行填充，使其横纵坐标的像素总数能被32整除，此处理的目的在于方便后续网络进行下采样。

将同批次的图像整体处理后，扩充第一维通道，在第一维通道上进行拼接。本实施例中，设置的batch size为2，则两张图片的输出为(2，3，W，H)。接着，将批次大小设置为2，利用数据集中的图像及标签，结合随机梯度下降方法对教师网络进行训练，训练的基础参数设置如下：学习率为0.01，动量设为0.9，权重衰减设为0.0001，训练过程中采用线性暖启动策略，迭代数量设为500，在周期为8和11时对学习率进行减小，总共训练12个epoch。检测头的分类损失采用focal loss，其中β参数为2，α参数为0.25；回归损失为一范数损失。

除基础的训练参数外，还需要对不同的蒸馏损失进行加权，不同的蒸馏损失其重要程度有所不同，本实施例中总共包含有三大部分的蒸馏损失，分别是骨干网络蒸馏损失、neck部分蒸馏损失以及检测头蒸馏损失，其中，骨干网络蒸馏损失中包含有全局特征权重、局部特征权重以及带权特征权重三个部分。

本实施例中，随机梯度下降方法的损失函数为：

式中，loss(·)表示损失函数，x_i表示标签的真实分布，y_i表示模型的预测分布，n表示标签的数量，在分类任务中采用交叉熵表示，即：

回归任务中采用均方误差表示，即：

蒸馏过程的激活函数为：

式中，z_i为第i个节点的输出值，z_j为第j个节点的输出值，T为蒸馏温度，如果T为1，则该广义softmax函数退化为一般的softmax函数，本实施例中的T＝5。

骨干网络蒸馏损失公式为：

式中，α＝0.1，β＝0.05，γ＝1000，p(x_i)表示真实概率分布，q(x_i)表示模型预测的概率分布；

蒸馏的总体损失公式为：

loss_dis＝λ₁loss_{backbone_dis}+λ₂loss_{neck_dis}+λ₃loss_{head_dis}

式中，loss_{backbone_dis}、loss_{neck_dis}、loss_{head_dis}分别代表骨干网络蒸馏损失、neck部分蒸馏损失以及检测头蒸馏损失，λ₁、λ₂、λ₃用于平衡各个蒸馏损失之间的权重，λ₁＝1，λ₂＝0.05，λ₃＝0.3。

网络总体损失的计算公式为：

loss＝loss_det+η·loss_dis

式中，

是交叉熵损失函数，

是smooth L1函数，η表示蒸馏损失权重，λ是归一化权重，t_i是一个向量，表示在RPN训练阶段，anchor相对于gt预测的偏移量，

是与t_i维度相同的向量，表示在RPN训练阶段，anchor相对于gt实际的偏移量，p_i为anchor预测为目标的概率，

有物体时为1，没有物体时为0，N_cls是总的anchor数量，N_reg是feature map的size。

在训练过程中，如果一开始就直接进行蒸馏容易导致网络不稳定产生梯度爆炸，因此训练过程中，蒸馏损失权重也会进行调整，本实施例采用的蒸馏损失权重调整策略的计算公式为：

式中，cur_iter表示当前训练的迭代次数；max_iter为超参数，表示蒸馏损失调整的最大迭代次数。本实施例中，当epoch为2时，开始蒸馏权重保持为1。

训练后，冻结教师网络参数，然后再次将图像送入教师网络中获得输出，并将输出结果连同标签一起送入学生网络中，得到蒸馏损失和目标检测损失后进行反向传播，最小化损失函数得到最终蒸馏后的学生网络，即得到压缩后的目标检测模型。

试验表明，采用本发明所设计的方案，不仅大幅提高了检测模型最终的准确性，而且降低了模型的参数量及提高了模型的计算效率，更加适用于诸如个人手机、小型嵌入式设备等算力更小的边缘设备。因此，与现有技术相比，本发明具有突出的实质性特点和显著的进步。

上述实施例仅为本发明的优选实施方式，不应当用于限制本发明的保护范围，凡在本发明的主体设计思想和精神上作出的毫无实质意义的改动或润色，其所解决的技术问题仍然与本发明一致的，均应当包含在本发明的保护范围之内。

Claims (10)

1.一种基于知识蒸馏的目标检测压缩方法，其特征在于，包括以下步骤：

步骤S1、对所有图像进行放缩操作，并确保图像缩小后比例与原始比例一致；

步骤S2、对图像进行预处理，然后将其进行拼接；

步骤S3、利用数据集中的图像及标签，结合随机梯度下降方法对教师网络进行训练，训练过程中，针对教师网络中的骨干网络蒸馏损失、neck部分蒸馏损失以及检测头蒸馏损失进行加权，其中，骨干网络蒸馏损失中包含有全局特征权重、局部特征权重以及带权特征权重三个部分；

步骤S4、训练后，冻结教师网络参数，然后再次将图像送入教师网络中获得输出；

步骤S5、将输出结果连同标签一起送入学生网络中，得到蒸馏损失和目标检测损失后进行反向传播，最小化损失函数得到最终蒸馏后的学生网络，即得到压缩后的目标检测模型。

2.根据权利要求1所述的一种基于知识蒸馏的目标检测压缩方法，其特征在于，所述步骤S1具体为：将所有图像进行放缩操作，使其最长边不超过1333像素，最短边不低于800像素，在放缩的同时避免图片长宽比失真，确保图像缩小后比例与原始比例一致。

3.根据权利要求1所述的一种基于知识蒸馏的目标检测压缩方法，其特征在于，所述步骤S2包括以下步骤：

步骤S201、将图像随机翻转，翻转的概率控制在0.5；

步骤S202、对图像像素值进行正则化处理；

步骤S203、将放缩后的图像进行填充，使其横纵坐标的像素总数能被32整除；

步骤S204、将同批次的图像整体处理后，扩充第一维通道，在第一维通道上进行拼接。

4.根据权利要求3所述的一种基于知识蒸馏的目标检测压缩方法，其特征在于，所述步骤S202中采用的正则化的计算公式为：

像素输出＝(像素值-均值)/方差

其中，三通道的均值和方差分别是[123.675,116.28,103.53]及[58.395,57.12,57.375]，该数值是由COCO训练集整体计算所得。

5.根据权利要求1～4任一项所述的一种基于知识蒸馏的目标检测压缩方法，其特征在于，所述步骤S3中，训练的参数设置如下：

学习率为0.01，动量设为0.9，权重衰减设为0.0001，训练过程中采用线性暖启动策略，迭代数量设为500，在周期为8和11时对学习率进行减小，总共训练12个epoch；

检测头的分类损失采用focal loss，其中β参数为2，α参数为0.25；回归损失为一范数损失。

6.根据权利要求5所述的一种基于知识蒸馏的目标检测压缩方法，其特征在于，所述步骤S3中，随机梯度下降方法的损失函数为：

式中，loss(·)表示损失函数，x_i表示标签的真实分布，y_i表示模型的预测分布，n表示标签的数量，在分类任务中采用交叉熵表示，即：

回归任务中采用均方误差表示，即：

7.根据权利要求6所述的一种基于知识蒸馏的目标检测压缩方法，其特征在于，所述步骤S3中，蒸馏过程的激活函数为：

式中，z_i为第i个节点的输出值，z_j为第j个节点的输出值，T为蒸馏温度。

8.根据权利要求7所述的一种基于知识蒸馏的目标检测压缩方法，其特征在于，所述步骤S3中，骨干网络蒸馏损失公式为：

式中，α＝0.1，β＝0.05，γ＝1000，p(x_i)表示真实概率分布，q(x_i)表示模型预测的概率分布；

蒸馏的总体损失公式为：

loss_dis＝λ₁loss_{backbone_dis}+λ₂loss_{neck_dis}+λ₃loss_{head_dis}

式中，loss_{backbone_dis}、loss_{neck_dis}、loss_{head_dis}分别代表骨干网络蒸馏损失、neck部分蒸馏损失以及检测头蒸馏损失，λ₁、λ₂、λ₃用于平衡各个蒸馏损失之间的权重，λ₁＝1，λ₂＝0.05，λ₃＝0.3。

9.根据权利要求8所述的一种基于知识蒸馏的目标检测压缩方法，其特征在于，所述步骤S3中，网络总体损失的计算公式为：

loss＝loss_det+η·loss_dis

式中，

是交叉熵损失函数，

是smooth L1函数，η表示蒸馏损失权重，λ是归一化权重，t_i是一个向量，表示在RPN训练阶段，anchor相对于gt预测的偏移量，

是与t_i维度相同的向量，表示在RPN训练阶段，anchor相对于gt实际的偏移量，p_i为anchor预测为目标的概率，

有物体时为1，没有物体时为0，N_cls是总的anchor数量，N_reg是featuremap的size。

10.根据权利要求9所述的一种基于知识蒸馏的目标检测压缩方法，其特征在于，蒸馏损失权重调整策略的计算公式为：

式中，cur_iter表示当前训练的迭代次数；max_iter为超参数，表示蒸馏损失调整的最大迭代次数。

Images