基于改進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%,方便部署在實際應用場景中。