基于改进YOLOv5s的轻量化安检图像检测算法研究①

马新月, 汤文兵

(安徽理工大学计算机科学与工程学院, 安徽 淮南 232001)

0 引 言

随着国家经济的飞速发展,公共基础设施建设也飞速发展,尤其是高铁、飞机、地铁等交通工具成为人们出行的首选。当面对大规模的人员流动时,安检人员需要在短时间内根据先验知识准确识别出刀具等危险品的形状,人工检测过程不仅要长时间保持高度集中注意力,还需要承担着如果没有准确检测出危险品将会给乘客带来危险的压力。因此,实现X光安检图像自动检测,是目前安检领域重要的发展趋势之一。

近年来,越来越多的深度学习方法被应用于安检图像违禁品检测。张友康[1]等人通过改进SSD[5]模型提出了适用于X光安检图像多尺度违禁品检测网络-ACMNet。该检测网络分为三个模块:通过小卷积对称模块增加细节特征,提高小尺度违禁品的检测;空洞多视野卷积模块通过上下文信息解决遮挡问题;多尺度特征融合策略在背景干扰的情况下,融合高、低层特征图来提高违禁品的检测精度。该算法获得了91.44%的mAP。穆思奇等[2]针对X射线图像中违禁品的自动检测问题,提出改进的YOLOv4[6]检测算法,利用空洞密集卷积模块增强特征表达能力和卷积视野,引用注意力机制,该算法在公开SIXray[7]数据集上平均精度达到80.16%。顾健等[3]在公开SIXray[7]数据集上进行图像处理,利用改进的Canny边缘检测方法去除背景噪声,对数据集进行有效的数据增强。在YOLOv3[8]的基础上加入SPP模块和一个新的anchor,来提高小目标违禁品的检测精度。文献[4]在单发多框检测器SSD的基础上,构建特征金字塔,引入可变形卷积获取危险品的形状和尺寸信息,得到一种X光危险品检测算法,来专门解决安检图像中小尺寸危险品检测效果差的问题。

以上基于卷积神经网络的方法检测多类别物体且存物体之间在相互遮挡时,准确率不高,且网络模型较大、检测速度较低。为了解决上述问题,实现X光安检图像地自动检测,提出一种改进的轻量化YOLOv5s的X光安检图像检测方法,称为EGD-YOLOv5 X光安检图像检测算法。该方法的优势主要在于:1)使用Ghost卷积层代替普通卷积层,减少了相同特征下的参数数量,减小了模型的大小,增加了原始图像中的信息量;2)通过添加高效通道注意力模块,使信道间特征相互关联,锚帧描述更加准确,提高了整个网络的检测精度;3)使用DIOU-nms替换原有的IOU-nms,对于某些存在遮挡的重叠目标,有效提升检测精度。该算法针对多品种违禁品检测,在保证模型精度的情况下,有效降低违禁品检测模型对计算资源的占用。

1 YOLOv5网络结构

YOLOv5是一种典型的一阶段目标检测算法,具有速度快、灵活性高的特点,在多个领域取得不错的成果[11]。从模型大小和检测精度两个方面考虑,选择以YOLOv5s为基础进行改进。由YOLOv5的网络结构可知,主要分为input、主干网络、Neck、输出端四部分。

1.1 输入端

输入端包含数据增强、自适应锚框计算、自适应图片缩放三种预先处理操作。在 YOLOv5 中除了使用最基本的数据增强方法外,还使用了 Mosaic 数据增强方法。Mosaic数据增强包括4-mosaic数据增强和9-mosaic数据增强两种,两者的区别就是,一个采用四张图片进行随机拼接,另一个采用九张图片进行随机裁剪、缩放以及拼接。自适应锚框计算是指YOLOv5能够根据不同数据集自适应的计算出不同的初始锚框比例,然后输出预测框,与真实框进行比较,反复迭代更新参数,直至找到最适合的锚框值。自适应图片缩放是指数据集中的图片各不相同,在输入网络之前会将其统一缩放到一个标准尺寸,即需要在图片周围填充不同大小的黑边,为了减少信息冗余,提高推理速度,YOLOv5对原始图像自适应的添加最少的黑边。

1.2 主干网络

Backbone,主干网络是用来提取特征的网络,作用就是提取图片中的信息,供后续的网络使用。主要有CBS,C3和SPPF模块组成,C3模块是由BottleneckCSP模块转变过来的,两者的作用基本相同。它由几个分支组成,每个分支包含一个卷积层和一个激活函数,所有分支的输出连接在一起形成 C3 输出。SPPF结构是YOLOv5作者基于SPP提出的,速度较SPP快很多,所以称为SPP-Fast。具体来说,就是用2个5*5的卷积操作代替SPP中9*9的卷积操作,3个5*5的卷积操作代替SPP中13*13的卷积操作。相同的感受野,但是计算的速度提升很多。

1.3 Neck模块

Neck,该模块的作用是对主干网络输出的特征图进一步加工,即特征信息的加强和特征融合,继续传递给Head。其网络结构设计也是沿用了FPN+PANet的特征金字塔结构,FPN结构将较强的语义特征从顶部特征映射传递到较低的特征映射。此外,PAN结构将较低的特征图传递到较高的特征图中。这两种结构共同增强了颈部网络的特征融合能力。

1.4 输出端

Head,为输出端,使用CloU Loss作为包围盒的损失函数,解决了标记盒与预测盒之间不重叠的问题;采用非最大抑制(Non-Maximum Suppression, NMS)机制增强了多目标和模糊目标的识别能力。

2 模型的改进与优化

虽然YOLOv5专注于目标检测,但是在多类别物体且物体之间存在相互遮挡时,平均准确率较低以及该网络模型较大不宜工业部署。为了解决上述问题,提出EGD-YOLOv5 X光安检图像检测算法,网络模型如图1所示,该模型的主要改进为:在YOLOv5s网络中引入Ghost模块和C3Ghost,代替普通卷积模块以及C3模块,以减少特征信道融合过程中的浮点运算。在主干网络中引入注意力模块(ECA),选择违禁品检测任务的关键信息,抑制非关键信息,从而提高算法的检测精度。使用DIOU_nms替换原有的IOU_nms,对于某些存在遮挡的重叠目标,DIOU_nms可以更好地保留正确的的检测框,从而提升检测精度。

2.1 轻量化网络模块

由于内存和计算资源的限制,在嵌入式设备上部署卷积神经网络(cnn)非常困难。在那些成功的卷积神经网络的设计中,特征图的冗余性很少被研究。但是特征层中充足或者冗余的信息总是可以保证对输入数据的全面理解,也可能是一个成功模型的重要组成部分,文献[9]设计轻量化模型时并没有试图去除这些冗余,而是用更低成本的计算量来获取它们。在特征层之间有很多是相似的,这些相似的特征层就像彼此的ghost(幻象)。该方法基于一种全新的神经网络基本单元--Ghost卷积模块,如图2所示。Ghost模块是对于输入的特征层,先用普通卷积操作生成部分真实的特征层(真实特征层),对于剩余的特征层(幻象特征层)通过对真实特征层进行线性操作获得,然后将真实特征层与幻象特征层拼接到一起组成完整的特征层。Ghost Bottleneck是一种类似于ResNet中的基本残差块,由两个堆叠的Ghost模块组成。第一个Ghost模块作为扩展层增加通道的数量,输出通道数与输入通道数之比称为扩展比。第二个 Ghost 模块减少通道数以适应shortcut路径,然后使用shortcut连接两个 Ghost 模块的输入和输出。Ghost Bottleneck中第二个Ghost module不再使用Relu,其他层在每层之后都应用了BN和Relu激活函数。C3Ghost是由Ghost Bottleneck模块来替换C3模块中的瓶颈模块,从而产生的一个新的模块。该结构取代了Bottleneck模块中的3*3标准卷积,降低了计算成本,压缩了模型尺寸。

图1 EGD-YOLOv5网络结构

图2 Ghost模块、GhostBottleneck模块和C3Ghost模块

2.2 注意力机制模块

在当前计算机算力资源的限制下,注意力机制绝对是提高效率的必要手段。将注意力集中到局部信息上,聚焦重要信息,忽略不相关信息。目前主要的注意力机制有:SE,CBAM,CA,ECA等。

虽然Ghost模块能够大幅度降低YOLOv5s模型的参数量和计算量,但不可避免地导致检测精度下降,为了提升检测精度,该方法在主干网络模块引入Wang[10]等人在2020年提出的一种高效通道注意力模块ECA(Efficient Channel Attention)。该模块避免了SENet中采用的降维操作对通道注意力带来的负面影响,以及有效地捕获了跨通道交互。该模块首先是对输入特征图进行空间特征压缩,即使用全局平均池化GAP,得到1*1*C的特征图;然后是通道特征学习,利用1D卷积得到等大小的特征向量;最后是通道注意力结合,将上一步得到的1*1*C的特征图与原始输入的特征图进行逐通道全乘操作,最终得到具有通道注意力信息的特征图。

在YOLOv5s骨干网络中,ECA模块嵌入在SPPF模块之前,ECA模块通过避免通道降维来学习有效的通道注意力,同时以极轻量级的方式捕获跨通道交互。

2.3 损失函数模块

损失函数的作用是测量预测神经网络信息与预期信息(标签)之间的距离,损失函数的计算公式如下:

L=Lbox+Lobj+Lcls

(1)

式中:Lbox采用CIOU函数,CIOU计算公式如下:

(2)

其中:IOU表示预测框和重叠框的交并比,ρ2(b,bgt)代表预测框和真实框的中心点之间的欧氏距离,α是权重参数,v是衡量宽高比。

NMS是大多数对象检测算法的最后一步,只要它与最高分框的重叠超过阈值,就会删除冗余检测框。在原始NMS中,当物体较近时,IOU值较大,经过NMS后处理,容易被抑制发生漏检的情况。在CIOU_Loss的基础上结合 DIOU_nms,对于某些遮挡重叠的目标,DIOU_nms可以更好地保留正确的检测框,因为抑制准则不仅考虑了重叠面积,还考虑了两个块之间中心点的距离。

3 结果和分析

3.1 试验数据集

试验采用百度飞浆的人工智能学习平台-X光安检图像挑战赛2022数据集,该数据集共有八个类别:刀(knife)、剪刀(scissors)、打火机(lighter)、优盘(USBFlashDisk)、压力容器(pressure)、带喷嘴塑料瓶(plasticBottleWithaNozzle)、公章(seal)、电池(battery),共计5017张图片,将5017张图像按照8:2的比例随机划分为训练集和验证机,其中训练集4014张、验证集1003张。

3.2 评价指标

为了验证改进后的YOLOv5s的性能,以精确度,mAP@0.5,参数数量,GFLOPs,FPS和模型大小作为衡量指标。其中mAP@0.5为IoU设置为0.5时所有类别AP的平均值,计算公式如下:

(3)

(4)

某个类别(C)的平均精度=验证集中该类(C)的所有精确率的和/含有该类别(C)目标的图像数量。

FPS用来衡量检测速度,指检测网络每秒可以处理的帧数,FPS 越高,网络模型处理帧的速度就越快。参数量和GFLOPs是用来衡量模型的网络复杂度,模型大小是指最终训练后保存的模型大小。

3.3 实验平台

实验的软硬件平台以及参数设置如下,GPU:RTX 3060,配置环境:CUDA 11.3,PyTorch 1.11.0,Python 3.8;batch_size 为 32;初始学习率为 0.01;总训练轮数为 300 轮。

3.4 结果与分析

分别使用默认参数对改进前和改进后的YOLOv5s网络进行了训练。改进后的损失变化如图4所示。

3.4.1 消融实验

为了分析各模块对模型的影响,设计消融实验,训练采用统一超参数,共验证5组实验,结果如表1所示,在YOLOv5中加入Ghost模块,模型的参数量从7.2下降到3.6,计算量下降了50%,模型体积压缩了45.8%,mAP下降了4.6%。实验1表明,加入ECA注意力模块后模型的精度有所提升,mAP指标比YOLOv5s提升了1.8%,并且几乎没有带来额外的参数量和计算量。实验3表明,在Ghost的基础上加入ECA模块,抵消使用C3Ghost带来的特征图通道间相关性不足,相对于实验2,mAP@0.5提升2.8%,参数量仅增加0.1。实验4表明,在实验1的基础之上将IOU_nms替换为DIOU_nms,算法的mAP@0.5提升0.7%。Ghost的引入可以有效降低网络复杂度,显著减少参数数量,但精度进一步降低,为了实现精度和速度的平衡,引入ECA和DIOU_nms来平衡因降低网络复杂度带来的精度损失,最终改进后的模型相比YOLOv5s,参数量减少了47.2%,计算量降低了46.4%,模型体积缩小了45.8%,FPS提升到99,mAP@0.5下降了0.7%。

图4 EGD-YOLOv5 Loss曲线

图5 不同算法结果对比

表1 消融实验

3.4.2 主流轻量化检测模型性能对比

为了验证该算法的轻量化效果,对比相同Baseline的轻量化网络YOLOv5s+MobileNet3,如表2,虽然YOLOv5s+MobileNet3网络的参数量仅有1.3,计算量2.3,以及模型体积仅有3.1,但是mAP@0.5仅有66.5%,所以综合mAP以及模型复杂度多方面考虑,该算法更具有优势,适合作为实际工业需求中违禁品的检测与识别。为了更直观地反映三种算法的性能差异,随机抽取5张图像,将3个模型放在同一个测试平台上进行测试,算法结果如图5所示。在图像2中,可以观察到,EGD-YOLOv5的检测效果优于YOLOv5s。

表2 轻量化检测模型性能对比

4 结 语

基于YOLOv5的改进的轻量化违禁品检测方法实现了X光安检图像的自动检测。首先针对遮挡情况,在CIOU_Loss的基础上采用DIOU_nms,然后采用ECA注意力机制对YOLOv5的特征提取网络进行改进,最后,利用Ghost模块轻量化的优势,引入Ghost BottleNeck模块对YOLOv5的C3模块进行重构,达到YOLOv5轻量化改进的效果。实验结果表明,EGD-YOLOv5的mAP达到93.6%,模型体积为7.8MB,浮点型计算量为8.9GFLOPs,与改进前的YOLOv5s网络相比,模型体积压缩了45.8%,计算量减少了46.4%,FPS提升到99,精度仅下降了0.7%,方便部署在实际应用场景中。