基于模型剪枝和知识蒸馏的船舶目标检测方法

邱 春,路 红,秦彬鑫,万文明,戴云周

(南京工程学院机械工程学院, 江苏 南京 211167)

近年来,随着目标检测的发展,对湖面和海面的船舶目标检测技术研究越来越受到重视,其在军事与民用方面都具有重要意义[1].目前,基于深度学习的目标检测技术需要大量的计算,卷积网络里大量的参数占用内存空间.为了解决神经网络中计算量大、占用内存多等问题,本文采用剪枝技术来对模型进行压缩[2],并用知识蒸馏技术对剪枝后的模型进行蒸馏,补偿剪枝后模型的检测精度损失.

1 模型选取

1.1 卷积神经网络

卷积神经网络常用于分类与检测等任务[3],主要由卷积层、池化层、全连接层组成.卷积层是通过感受野即卷积核在图像上进行滑动卷积,以此来提取图像中不同维度的特征;将卷积层输出的特征放到池化层中对特征图进行池化操作,池化层与卷积层相似,是在感受野上建立一个窗口来进行滑动操作,从而达到特征降维的作用;将池化后输出的特征送到全连接层中进行维度的转换,在全连接层中对特征进行分类等操作,输出不同类型的特征.典型的卷积神经网络有VGG、AlexNet、ResNet等.

1.2 YOLO v3算法介绍

在计算机视觉领域中,基于深度学习的目标检测算法主要分为两阶段算法和单阶段算法:1) 两阶段算法有R-CNN、Faster R-CNN算法,先获得候选区域(region proposal),然后将其分成不同的目标类别,该类算法的定位和检测准确率较高,但运算速度较慢,实时性达不到要求;2) 单阶段算法有YOLO、SSD、YOLO900算法等,依据回归思想使实时检测成为可能,但检测精度比两阶段算法低.本研究模型选择单阶段算法中的YOLO v3算法.YOLO v3算法主要由特征提取网络Darknet53和检测网络组成,Darknet53网络由53个卷积层组成,检测网络是在不同尺度的特征图上选取3个尺度的特征图进行目标预测输出[4].YOLO v3算法检测原理是:首先将输入的图像送到Darknet53网络中进行特征提取,得到多个尺度不一的特征图,对特征图进行处理获得目标的相关类别与回归框;然后经过极大值抑制与类别置信度对目标边框进行过滤;最后得到合格的目标信息.YOLO v3算法检测精度高、速度快,满足了人们对目标检测的实时性要求.YOLO v3算法的网络结构图如图1所示.

图1 YOLO v3算法网络结构图

2 船舶检测模型设计

2.1 模型剪枝

基于深度学习的模型极其复杂,实际应用中往往会受到计算资源消耗大、存储空间要求高等阻碍,对于嵌入到微型机或者移动端来说更加困难.神经网络模型中有大量的神经元与权重冗余,真正对模型产生影响的只有5%～10%.将模型中不重要的参数去掉来减小模型大小的方法越来越多,最常见的模型压缩方法是模型剪枝[5-6].剪枝不仅可以降低模型复杂度、防止过拟合,还可以压缩模型大小、提高模型的速度.剪枝方式可分为通道剪枝、权值剪枝、层剪枝.层剪枝效果最粗糙,但无软硬件要求、灵活性好;权值剪枝灵活性最好,但是对软硬件有要求,需要专门的系数矩阵运算库或者硬件来做加速;通道剪枝集成了层剪枝和权值剪枝的优点,精细且灵活性好.本研究采用通道与层相结合方式来对模型进行剪枝.模型剪枝示意图如图2所示.

图2 模型剪枝示意图

通道剪枝对软硬件要求不高,在细粒度方面进行了很好的平衡,采用过随机抛弃channel-wise的连接引入稀疏化,将常见的L1约束添加到各个通道的缩放因子上,使其能够更加便捷地优化目标,通道剪枝的目标函数为:

(1)

式中:W为训练的权值;x、y为训练数据中的输入、输出;γ为批量归一化(BN)的缩放因子;g(γ)为稀疏对比例因子的惩罚函数;λ为平衡因子.

通道剪枝主要是引入BN中的缩放因子γ和稀疏训练中的惩罚项L1.最常见的操作方式是首先在模型中的每个通道都放入一个γ,与通道的输出结果进行乘积运算,在进行稀疏化时将这些比例因子与联合训练的网络权重结合起来进行正则化；然后按照缩放因子的大小进行排序,设置一个阈值,根据阈值将小的缩放因子通道直接裁剪掉,重复多次进行该操作,选取效果较好的一组；最后对网络进行微调即可[7].

层剪枝操作是针对模型中每一层的shortcut层的前一个CBL进行评价,每层γ的取值按照从大到小规则进行排序,取最小的进行层剪枝,每剪去一个shortcut结构,相当于减掉1个shortcut层和它前面的2个卷积层[8].

在通道剪枝结束后,将模型再汇总进行层剪枝,最后输出剪枝完的模型.

2.2 知识蒸馏

目前,深度学习需要对海量的数据进行训练来提取特征,这些数据数量较大且冗余,需要消耗大量资源,在没有实时性要求下训练得到的模型精度较高,但是模型结构笨重.在实际应用中,考虑各方面成本,特别是在计算资源方面有所限制,使用一个较大的神经网络模型有较多局限性.因此,本文采用知识蒸馏模型,让大模型为小模型进行知识迁移.

知识蒸馏的原理是将复杂的教师网络中的dark knowledge迁移到简单的学生网络中[9],知识蒸馏模型示意图如图3所示.

图3 知识蒸馏模型示意图

知识蒸馏模型用简单目标来辅助困难目标一起训练,soft target来自于教师网络的预测输出[10],通过在softmax中增加温度参数T来实现简单目标,公式为:

(2)

式中:qi为函数的目标输出;T为温度参数,是从统计力学的玻尔兹曼分布中借用的概念,当温度T趋向于0时,softmax输出将收敛为一个one-hot向量,温度T趋向于无穷时,softmax的输出则更「软」;Z为softmax层输出的分类类别概率.

在训练新模型时,使用较高的T使softmax产生的分布足够软,让新模型(同样温度下)的softmax输出近似原模型;在训练结束后使用正常的温度T来预测.

3 试验与分析

本文所有的模型训练及相关试验均在深度学习框架pytorch下实现,试验的硬件配置为:Intel i5处理器,8 G内存、NVIDIA GTX 2070super GPU、 PC、Windows 10 操作系统.

3.1 数据集的制作

国际上公开的海面船舶目标检测数据集较少,为评估本文船舶目标检测模型的准确性,自建船舶目标检测的小样本数据集.通过网上截取船舶视频片段进行视频的分帧得到一定数量的船舶图像,并采用几何变换中的旋转、裁剪、像素变换中的高斯噪声等方式进行数据集的扩增,即先对数据集采用旋转方式进行扩增,再随机选取一部分图像进行裁切和高斯噪声操作,将两次扩增的数据集集合构成总的船舶数据集;采用LabelImg标定工具对数据集进行标定,得到符合YOLO格式的txt标签文件;将标注后的图像数据集中的样本按照7∶2∶1的比例划分为训练集、验证集和测试集,扩增后的船舶数据集信息如表1所示.

表1 船舶数据集信息

3.2 试验与结果

首先对船舶目标检测模型进行剪枝压缩,将制作的船舶数据集放入YOLO v3模型中进行训练；然后对模型进行稀疏训练后进行通道剪枝、层剪枝；最后对剪枝后的模型进行微调,在微调的过程中对模型进行知识蒸馏,使模型得到更好的效果.

表2为不同剪枝系数的通道剪枝结果对比.由表2可见,当剪枝系数在0.85时效果最好,经过剪枝后模型的参数量由61 523 734个减少到2 323 692个,减掉约96%,检测精度由0.931 304降低到0.920 625,下降近1%,总体效果较好.

表2 不同剪枝系数的通道剪枝结果对比

通道剪枝后进行层剪枝,表3为不同剪枝系数的层剪枝结果对比.由表3可见,剪枝系数增大，检测精度不断下降,参数量下降缓慢,层剪枝没有通道剪枝效果明显；剪枝系数为3时检测精度最高,与原始检测精度接近,参数量没有太多减少；综合来看，选取剪枝系数为3来进行层剪枝效果最好.

表3 不同剪枝系数的层剪枝结果对比

图4为模型剪枝前后的检测效果对比图.由图4可见,剪枝后的模型检测精度与剪枝前相比没有太多下降,甚至在小目标检测中剪枝后的模型检测精度比剪枝前有所提升.

(a) 剪枝前

(b) 剪枝后

(c) 剪枝与微调后

4 结语

本文采用模型剪枝与知识蒸馏技术对船舶目标检测模型进行压缩.通过剪枝技术将模型的参数量减少了超过90%,大大压缩了模型体积,使该模型能够部署到内存空间条件不足的硬件设施上,提升了模型的计算速度,提高了模型的实时性,解决了大部分硬件设施运行慢的问题.对模型进行剪枝后采用知识蒸馏技术对检测精度进行微调,提高了模型的检测精度.