基于模型压缩YOLOv4的弹载图像舰船目标实时检测

雷 鸣,王曙光,凌 冲,杨传栋,秦 杰

(陆军炮兵防空兵学院 高过载弹药制导控制与信息感知实验室, 合肥 230031)

摘要:根据深度学习目标检测算法YOLOv4弹载平台移植的实时性和轻量化需求,针对YOLOv4网络模型参数冗余导致的储存空间成本高和运算资源占用大等问题,在PC端使用舰船数据集训练模型并对其进行模型压缩。采用L1正则化稀疏训练的方法对预训练模型进行BN层Gamma系数稀疏化,通过剪枝技术在精度下降较少的情况下减少模型参数并压缩模型结构,使用知识蒸馏技术使模型精度进一步回升。实验结果表明,压缩后模型体积减少92%,参数减少91.9%,推理速度提升2倍,实现了高精度实时检测。

关键词:目标检测;深度学习;弹载图像;舰船目标;模型压缩;知识蒸馏

1 引言

如何将性能优越的目标检测算法移植到弹载平台,实现图像制导弹药的自寻的打击,并且满足准确率、功耗和实时性要求是图像制导弹药的重要发展趋势。弹载平台是指通过舰炮、榴弹炮、迫击炮和火箭炮等发射的精确制导弹药上搭载的嵌入式平台。弹载平台是一个高速旋转且不断前进的平台,会造成图像抖动、旋转、失焦和尺度变化较大等问题,不仅对目标检测算法的实时性和精度有着较高要求,而且需要算法有较好的尺度适应性和鲁棒性。弹载图像目标检测技术指的是利用弹载摄像机获取的图像进行逐帧(场)分析计算,获得目标的位置、大小和类别信息的技术,是智能弹药发展的关键技术之一。传统弹载图像目标检测算法主要有模板匹配法、特征统计法、图像分割法等。虽然传统目标检测算法在硬件平台上实际应用成熟,但是人工设计算法来提取特征对研究者经验要求较高,在光照、尺度和背景变化的情况下鲁棒性不足。近年来,基于深度学习的目标检测技术得到广泛应用,省去了人工提取特征的步骤,是一种更为高效的方法。

值得注意的是,深度学习算法在弹载平台上的移植要考虑运算资源占用的问题。目前,深度学习目标检测算法YOLOv4(You Only Look Once v4)[1]集成多种先进算法策略,在MS COCO数据集的实时检测速度约为65 FPS,平均精度均值为43.5%,相比上一代YOLOv3算法分别提高12%和10%。虽然YOLOv4算法在PC端利用GPU加速,可以做到实时高精度检测,但是其权重文件达到了244 MB,在运算资源有限且条件复杂的弹载平台上,检测启动慢、预存开销大且运算成本高,必须针对特定打击目标进行神经网络模型的针对性训练,并且对模型进行轻量化处理,在精度下降不大的情况下提高检测速度,以更好地适应弹载平台硬件的特殊需求。

2 YOLOv4目标检测算法

为实现效果最佳的实时高精度检测,本文选取一阶段深度学习目标检测算法YOLOv4作为基础网络。YOLOv4算法首先将输入图像分割成K×K的栅格网格,然后每个网格负责检测中心坐标位于网格内的目标,同时网格内会预测边框的位置坐标和边框的置信度,最后对预测边框进行筛选时使用非极大值抑制算法选出最优边框[2]。YOLOv4算法检测网络是由改进的输入端、基于CSP Darknet53的特征提取网络、SPP(Spatial Pyramid Pooling)附加模块、PANet(Path Aggregation Network)特征融合模块以及基于锚框机制的YOLOv3预测模块组成的体系结构。本算法网络结构如图1所示。

图1 YOLOv4网络结构框图

Darknet53特征提取网络分为5个大残差块,其中包含的小残差单元的个数分别为1、2、8、8、4。CSP Darknet53是在Darknet53的每个大残差分块上加上CSP(Cross Stage Partial)模块[3],CSP模块将上一层的特征映射划分为两部分,然后通过跨阶段层次结构将它们合并,608×608的特征图经过5个CSP模块后输出19×19的特征图。CSP模块可以增强神经网络的学习能力,能够在保证网络结构轻量化的同时减少内存成本、降低运算瓶颈和维持准确性。CSP残差块网络结构如图2所示。

图2 CSP残差块网络结构框图

其中,BN[4]指的是批量归一化(Batch Normalization),起着优化训练的作用,可缓解训练时的梯度消失与梯度爆炸问题,批量是指训练模型设置的图片数量。YOLO系列网络模型都引用了BN层,在每个卷积层后面添加,防止了过拟合且降低了模型向前推理速度。BN计算过程为

(1)

式中: γ为缩放参数; μ为均值; σ2为方差; β为偏置量;ω为权值参数; xconv为通道特征图卷积计算值。

卷积层和BN层合并后,权值参数变为:

(2)

偏置量变为:

(3)

最终,合并后的BN计算公式为:

(4)

卷积层和BN层的合并可以减少运算成本,有利于提升模型速度与模型泛化能力。输入图片经过卷积层得到特征图,接着经过BN层的归一化处理后由激活函数激活。其中,γβ是用来缩放和偏置的2个参数,通过神经网络训练迭代得到,体现每层网络的特征分布。本文剪枝过程基于BN层γ缩放参数实现[5],不需额外引入其他参数,可大幅减小模型参数和计算资源占用量,从而提升神经网络模型的推理速度。

3 模型压缩策略

神经网络模型中有大量冗余参数,其存在对目标检测结果无影响,这是因为神经网络模型在训练过程中需要保有一个大参数量来寻找最优解[6]。这些大部分无用的参数在训练结束后还是保留在了权重文件中,使得权重文件过大。本文模型压缩策略通俗来讲就是在原始参数范围内找到对检测效果影响最大的计算路径,剪枝掉冗余的通道和层,这样神经网络模型的精度不会降低且占用储存空间大幅减少。模型剪枝的策略有细粒度剪枝[7]和粗粒度剪枝[8]两种类型,分别对应权重剪枝和通道剪枝。细粒度剪枝在目标类别数量多的情况下效果较好,本文研究的是针对舰船目标检测的模型剪枝,检测目标类别较少。粗粒度剪枝具有原理简单、计算开销小、硬件依赖小和可行性高等特点,针对特定数据集的剪枝效果较好,所以本文采用粗粒度的通道剪枝对神经网络模型进行压缩。

本文模型压缩的策略是:首先,在PC端使用YOLOv4训练舰船目标数据集,获取权重文件为后续工作做准备;其次,对神经网络模型进行稀疏训练;再次,对神经网络中冗余的通道或层进行剪枝,压缩神经网络模型的深度和宽度;最后,对剪枝后的模型进行微调来回升精度。

3.1 稀疏训练

缩放参数γ决定着BN层特征图通道的重要性,代表着每一通道的权重。假设输入通道i对应的γi值接近0,那么式(1)分子部分随着γi的减小而减小从而接近常数0,导致此通道卷积计算值对后续计算不产生影响[9]。通常来讲,训练后的神经网络中BN权重一般呈正态分布,等于或者接近0的参数较少。显然,如果想要完成冗余参数的剪枝,那么就要在神经网络训练时减小缩放参数γ值,这种方法叫做稀疏训练。引入稀疏规则算子L1范数(Lasso Regularization)进行正则化,计算过程为

(5)

式中,前项对应于卷积神经网络的正常训练损失,其中(xy)表示训练的输入和目标边框,W表示可训练的权重,后项中的g(γ)是对缩放因子的稀疏性惩罚,s是用于平衡损失和稀疏性惩罚的超参数,引入L1范数g(γ)=|γ|,对于非平滑L1惩罚项,采用次梯度作为优化方法。 另一种选择是用平滑L1罚分代替L1罚分[10],以避免在非平滑点使用次梯度。

超参数s根据数据集大小、平均精确率mAP(mean Average Precision)和BN权重分布进行调整,稀疏训练中s越大稀疏就越快,但mAP下降的也越快。学习率调整更大,稀疏会加快,但是小的学习率有助于后期mAP的回升。通过调整学习率μs来保持mAP下降较少而同时实现模型的高稀疏度。

稀疏训练整个训练过程保持恒定的s,这样可以保持均匀的压缩力度从而获得较高的压缩度,但是此方法对精度的折损较大,且最优s值的调参过程时间成本较大。如果在训练进行一半的阶段把s值衰减100倍,有助于精度的回升,减少了寻找最优s值的时间成本,但是会牺牲一部分压缩度。

3.2 模型剪枝

模型剪枝属于模型压缩的范畴,实际应用成熟。通道剪枝即有选择地剪枝删除某个卷积核对应通道。层剪枝即有选择地直接删除某个卷积层。

剪枝效果取决于稀疏训练的效果[11],BN权重稀疏度越高,粗粒度剪枝后模型压缩度越大。通道剪枝后可以进行层剪枝,进一步提升压缩度。本文通道剪枝利用全局阈值筛选出各卷积层的掩膜,然后和与之相连的直接连接层(shortcut)取并集进行剪枝,此方法考虑了每一个相关卷积层,并且限制了每一卷积层的保留通道。层剪枝的原理与通道剪枝的原理相似,本文通过对每一个直接连接层前的残差块进行打分,排列各层γ值分数,对最小得分的层进行剪枝来压缩骨干网络。为保证神经网络模型的完整性,对直接连接层做处理时同时剪掉其之前的2个卷积层。

3.3 知识蒸馏

2015年Hinton团队首次提出知识蒸馏(Knowledge of distillation,KD) 的思想[12]。针对剪枝导致的精度下降问题,本文对剪枝后的神经网络模型进行微调以回升精度。剪枝后的小模型与原始模型在结构上相似,可采用知识蒸馏的策略进行微调。知识蒸馏是基于TS(teacher-student) 框架的模型压缩策略,指一个参数量大的复杂神经网络模型(teacher)学习到的特征表征能力传递给一个参数量小的简单神经网络模型(student),从而获得一个速度快、参数少且学习能力强的模型。具体算法流程如图3所示。

图3 知识蒸馏流程框图

利用标注目标边框的数据集训练大模型,即基础训练模型,设定一个蒸馏温度值(Temperature,T),然后用基础训练模型的预测输出值除于T,再进行soft max回归[13],由此获得介于0~1的软目标(soft target)预测概率值。然后将基础训练模型的特征表征能力传递给小模型,这一阶段分为两部分进行,一是一部分小模型预测输出除以T然后进行softmax回归,二是剩余部分小模型预测输出直接进行softmax回归。硬目标(hard target)指的是真实目标,0表示无真实目标,1表示有真实目标。参数λ是权重参数,用于调节基础训练模型与小模型之间交叉熵损失函数的权重,作为小模型训练的损失函数之一。

4 实验结果与分析

本研究实验环境为Microsoft Windows 10专业版(64位) 操作系统,英特尔Intel(R) Core(TM) i5-8300H @ 2.30GHz(2304 MHz) CPU,HP 84DB主板,8.00 GB(2667 MHz) 内存,NVIDIA GeForce GTX 1060(6144MB) 显卡,CUDA版本10.1.168,cudnn版本7.6.4,OpenCV版本3.4.1,Visual Studio版本2015。基于pytorch框架实现模型压缩算法。

本研究采用A7200CG30彩色面阵工业相机拍摄水面上的舰船模型,模拟弹载摄像机获取图像时的光照条件、下滑角与尺度变化,同时考虑到了实际情况中目标出视场或处于视场边缘对目标检测效果的影响,拍摄了部分处于图像边缘的舰船目标图像,充分保证舰船目标数据的真实性与多样性。获取图片3 548张,随机选取2 838张图片作为训练集,选取355张图片作为验证集,选取剩余355张图片作为测试集,使用Yolo-mark工具进行目标边框标注。在训练过程引入马赛克(Mosaic)数据增强方法,把4张具有不同语义信息的输入图像通过随机裁剪、随机排列、随机缩放的方式拼接成一张图像,然后通过批量归一化的操作从每层的拼接图像中获取激活统计信息。如此可以使检测器检测超出常规语境的目标,进一步丰富了数据集,同时可以减少对微型批量(mini-batch)大小的依赖从而增强模型的鲁棒性,训练过程可视化如图4所示。

图4 Mosaic数据增强训练过程可视化图形

首先进行模型的基础训练,调整学习率等参数获得效果最佳的模型作为稀疏训练的预训练模型。设置稀疏训练超参数s,经过反复测试,通过可视化工具查看BN层中γ稀疏的稀疏情况,调整超参数s与学习率。根据先验知识,超参数s值设置区间为10-3~10-5,根据数据集、学习率、BN分布和mAP调整,学习率越小越有助于稀疏训练后期精度提升,s值越大γ系数的压缩越快,同时精度也下降越快。经过循环测试,最终设定s为10-4,在训练中段s衰减之前可以达到较高压缩度的同时,在s衰减之后恢复精度能力较好。稀疏训练前后γ系数的压缩程度如图5和图6所示,可以看到γ系数大部分都被压缩至接近0,γ系数接近0的通道影响作用极小,可以进行通道剪枝。

图5 稀疏训练前的γ系数分布图形

图6 稀疏训练后的γ系数分布图形

稀疏训练过程可视化如图7所示,可见稀疏训练在较小影响精度的情况下完成了γ系数压缩的任务,符合剪枝条件。稀疏训练前后各评价指标对比如表1所示。

表1 稀疏训练前后各评价指标

训练阶段查准率P召回率RmAPF1基础训练0.6010.8570.7610.683稀疏训练0.5860.9190.7480.75

图7 稀疏训练各评价指标曲线

YOLOv4神经网络原始模型有161层,其中有104 个BN层,通道数量为32 256个,包含63 943 071个参数,权重文件大小为244 MB。可见参数量与体积巨大,不适合在弹载平台部署,所以对其进行剪枝处理。

全局通道剪枝0.85,层保留0.01,剪枝后的神经网络模型参数为7 033 738,被压缩了89%,通道数为11 153,被压缩了65.4%,模型大小为25.7 MB,被压缩了89.5%。mAP为0.698,可以继续对冗余的卷积层进行剪枝。

YOLOv4中有23个直接连接层,对于舰船数据集这种各类别特征较为相似的简单数据集,可以在剪枝大量直接连接层的同时精度衰减较少,本研究剪枝了16个直接连接层和其前面的两个卷积层,共48层,层剪枝后BN层数为56层,减少了46.2%,总层数为113层模型大小19.9 MB,mAP为0.621。

为了恢复精度,使用知识蒸馏的思想对剪枝后模型进行微调。设置原始模型为老师模型,剪枝后模型为学生模型,使简单模型可以学习到复杂模型的特征表征能力,进而回升精度。

各模型效果如表2所示,A代表YOLOv4原始模型,B代表通道剪枝后模型,C代表层剪枝后模型,D代表知识蒸馏微调后的模型。

表2 各模型评价指标

模型mAPFPS通道数层数量参数量大小/MBA0.76123.032 2561606.4×107244.0B0.69855.811 1531607.0×10625.7C0.62168.55 4351125.2×10719.9D0.78266.95 4351125.2×10719.9

对舰船目标检测效果如图8所示。

图8 舰船目标检测效果图形

最终微调后的神经网络模型相较于原始模型精度稍有提升,这是因为稀疏训练和知识蒸馏训练也相当于进行了L1正则化,推理速度提升了2倍,通道数缩减为原始模型的16.8%,剪枝了46.2%的层,参数量规模下降91.9%,模型大小为原始模型的8%。

5 结论

1) 提出了一种可以完成实时高精度检测的轻量级舰船目标深度学习检测模型,对硬件设备要求更小,充分利用运算资源,为深度学习算法的弹载平台移植提供了一个可行预研方案。

2) 此模型主要针对水面环境的舰船类目标检测,对其他战场环境下的目标检测能力有待考证,接下来需进一步在弹载平台上进行验证。

参考文献:

[1] ALEXEY BOCHKOVSKIY,WANG C Y,HONG-YUAN MARK LIAO.YOLOv4:Optimal Speed and Accuracy of Object Detection[EB/OL].[2020-04-23].https://arxiv.org/abs/2004.10934v1.

[2] REDMON J,DIVVALA S,GIRSHICK R,et al.You onlylook once:unified,real-time object detection[C]//Proceedings of IEEE Conference on Computer Vision and Pattern Recognition.Washington D.C.,USA:IEEE Computer Society:2016:779-788.

[3] WANG CHIENYAO,HONG-YUAN MARK LIAO,W U YUEH-HUA,et al.CSPNet:A new backbone that can enhance learningcapability of cnn[C]//Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshop.[S.l.]:[s.n.],2020.

[4] 王岩,吴晓富.深度神经网络训练中适用于小批次的归一化算法[J].计算机科学,2019,46(S2):273-276,308.

[5] LIU Z,LI J,SHEN Z,et al.Learning Efficient Convolutional Networks through Network Slimming[C]//2017 IEEE International Conference on Computer Vision(ICCV).IEEE,2017:2755-2763.

[6] 雷杰,高鑫,宋杰,等.深度网络模型压缩综述[J].软件学报,2018,29(02):251-266.

[7] 昌晶.细粒度模型剪枝算法研究[D].南京:南京大学,2019.

[8] M.HAJABDOLLAHI,R.ESFANDIARPOOR,et al.Sim-plification of neural networks for skin lesion image segmentation using color channel pruning[J].Computerized Medical Imaging and Graphics,2020,82.

[9] 刘建伟,赵会丹,罗雄麟,等.深度学习批归一化及其相关算法研究进展[J].自动化学报,2020,46(06):1090-1120.

[10] SCHMIDT M,FUNG G,ROSALES R.Fast optimization methods for l1 regularization:A comparative study and two new approaches[J].In ECML,2007:286-297.

[11] 靳丽蕾.基于剪枝的卷积神经网络压缩方法研究[D].保定:河北大学,2019.

[12] HINTON G,VINYALS O,DEAN J.Distilling the Knowledge in a Neural Network[J].Computer ence,2015,14(7):38-39.

[13] TOM MIKOLOV,KOMBRINK S,LUK BURGET,et al.Extensions of recurrent neural network language model[C]//IEEE International Conference on Acoustics,Sp-eech & Signal Processing.IEEE,2011:5528-5531.

Real-Time Detection of Ship Target in Missile-Borne Platform Images Based on Model Compressed YOLOv4

LEI Ming, WANG Shuguang, LING Chong, YANG Chuandong, QIN Jie

(Laboratory of Guidance Control and Information Perception Technology of High Overload Projectiles, PLA Army Academy of Artillery and Air Defense, Hefei 230031, China)

Abstract: According to the real-time and lightweight requirements of the deep learning target detection algorithm YOLOv4 missile-borne platform transplantation, in order to solve the problems of high storage space cost and large computing resource occupation caused by the redundancy of the YOLOv4 network model parameters, the ship data set was used to train the model on the PC side and perform model compression on it. L1 regularization sparse training method was used to sparse the pre-trained model Gamma coefficients in the BN layer, and then the pruning technology was used to reduce the model parameters and compress the model structure with less accuracy degradation, and finally use the knowledge distillation technology to make the model more accurate pick up. Experimental results show that after compression, the model volume was reduced by 92%, the parameters were reduced by 91.9%, and the inference speed increased twice, realizing high-precision real-time detection.

Key words target detection; deep learning; missile-borne image; ship target; model compression; knowledge distillation

收稿日期:2020-11-10;修回日期: 2020-12-13

作者简介:雷鸣(1996—),男,硕士研究生。

doi: 10.11809/bqzbgcxb2021.09.036

本文引用格式:雷鸣,王曙光,凌冲,等.基于模型压缩YOLOv4的弹载图像舰船目标实时检测[J].兵器装备工程学报,2021,42(09):225-230.

Citation format:LEI Ming, WANG Shuguang, LING Chong, et al.Real-Time Detection of Ship Target in Missile-Borne Platform Images Based on Model Compressed YOLOv4[J].Journal of Ordnance Equipment Engineering,2021,42(09):225-230.

中图分类号:TP391

文献标识码:A

文章编号:2096-2304(2021)09-0225-06

科学编辑 杨继森 博士(重庆理工大学教授)责任编辑 杨梅梅