引入注意力機制的改進型YOLOv5網絡研究

曾 凱，李 響，陳宏君，文繼鋒

(南京南瑞繼保電氣有限公司，江蘇 南京 211102)

zengkai2@nrec.com;lix@nrec.com;chenhj@nrec.com;wenjf@nrec.com

1 引言(Introduction)

基于深度學習算法的目標檢測技術已在計算機視覺領域得到廣泛應用[1-3]，然而在面對復雜場景下的多尺度目標檢測時，仍然存在識別精度不夠高、定位不夠準確的問題。目前，大部分研究是在基準檢測模型上進行多尺度特征增強和引入混合注意力機制模塊。比如，林森等[4]提出了一種基于注意力機制與改進YOLOv5網絡的水下珍品檢測方法，其主要思想是利用CBAM注意力機制模塊對特征提取網絡進行改進；劉萬軍等[5]針對背景復雜的遙感圖像識別任務，在Faster R-CNN基準網絡上提出多尺度特征增強及密集連接網絡，有效地解決了目標漏檢的問題。綜合各方面的研究表明，引入注意力機制模塊可以有效地提升目標檢測的準確率。

注意力機制可分為基于卷積神經網絡的注意力機制[6]和基于Transformer網絡的自注意力機制[7]兩大類。注意力機制模塊眾多、模型性能差異大，對比評估一些新型且有效的注意力機制模塊，對提升復雜多尺度目標的檢測性能是非常有意義的。

本文選取YOLOv5網絡作為基準模型，在網絡中引入多種先進的注意力機制模塊，設計并實驗對比分析了幾種改進型網絡的性能表現。

2 改進型YOLOv5網絡的設計(Design of improved YOLOv5 network)

2.1 標準YOLOv5目標檢測網絡結構

YOLOv5網絡在工業落地應用中表現出極其優秀的檢測性能和推廣價值，它整合了大量的計算機視覺前沿技術，顯著改善了對象檢測的性能，提升了模型訓練的速度及模型應用的便利度。

YOLOv5網絡結構示意圖如圖1所示。

圖1 YOLOv5網絡結構示意圖Fig.1 Diagram of YOLOv5 network structure

圖1中的基準網絡主要由骨干特征提取網絡(Backbone網絡)、頸部特征融合網絡(Neck網絡)和檢測頭部預測網絡(Prediction網絡)組成，分別主要使用基于瓶頸(Bottleneck)結構的C3為Backbone，基于多特征圖融合的FPN+PANet結構為Neck，以及基于檢測目標的位置和類別進行回歸、分類任務的YOLO檢測頭為Prediction。

2.2 引入卷積神經網絡注意力機制的改進型YOLOv5網絡設計

2.2.1 引入協同注意力機制網絡的改進型YOLOv5網絡

協同注意力機制網絡(Coordinate Attention，CA)[8]，它引入了一種新的注意塊結構，該結構不僅能捕獲跨通道的信息，還能捕獲方向感知和位置感知的信息，這能幫助模型更加精準地定位和識別感興趣的目標。

基于CA模塊的典型網絡結構如圖2所示。

圖2 CA注意力模塊的典型網絡結構示意圖Fig.2 Network structure diagram of CA attention module

圖2中的網絡結構首先分別對水平方向和垂直方向進行全局平均池化，得到兩個1 維向量，在空間維度上拼接并經過1×1的卷積壓縮通道數，然后通過批量歸一化(BN)和非線性激活函數編碼垂直方向和水平方向的空間信息，接著在空間維度上將BN和激活函數的輸出拆分成兩個特征圖，再各自通過一個1×1的卷積調整通道得到和輸入特征圖一樣通道數的融合了注意力機制的特征圖。

本文將CA網絡模塊應用在YOLOv5網絡的骨干特征提取網絡的每個多尺度特征輸出的位置，用于對每個尺度下特征圖的各通道進行特征重標定，以提升原始YOLOv5網絡的特征提取能力。設計的引入CA網絡模塊的改進型YOLOv5網絡結構示意圖如圖3所示，命名為CA-YOLOv5網絡。

圖3 CA-YOLOv5網絡結構示意圖Fig.3 Diagram of CA-YOLOv5 network structure

2.2.2 引入CBAM網絡的改進型YOLOv5網絡

CBAM網絡在通道注意力網絡的基礎上擴展了一個空間注意力模塊，它可以在通道和空間維度上進行注意力運算。CBAM網絡[9]包含兩個子模塊：通道注意力模塊(Channel Attention Module，CAM)和空間注意力模塊(Spatial Attention Module，SAM)。

基于CBAM注意力模塊的典型網絡示意圖如圖4所示。

圖4 基于CBAM注意力模塊的典型網絡示意圖Fig.4 Network structure diagram of CBAM attention module

CAM模塊通過一個并行的最大值池化層，得到兩個1×1×C的特征圖，首先將得到的兩個特征向量相加，然后經過一個激活函數得到每個通道的權重系數，最后用權重系數與原來的特征圖通道相乘，即可得到縮放后的新特征圖。

SAM模塊先分別進行一個通道維度的平均池化和最大池化，得到兩個H×W×1的通道張量，將這兩個張量在通道維度進行拼接，再經過一個7×7的卷積及Sigmoid激活函數，得到權重系數，最后用權重系數和輸入特征圖相乘，即可得到縮放后的新特征圖。

本文將CBAM網絡模塊應用在YOLOv5網絡的骨干特征提取網絡的每個多尺度特征輸出的位置，用于對每個尺度下特征圖的各通道/各空間進行特征重標定，以提升原始YOLOv5網絡的特征提取能力。設計引入CBAM網絡的改進型YOLOv5網絡結構示意圖如圖5所示，命名為CBAMYOLOv5網絡。

圖5 CBAM-YOLOv5網絡結構示意圖Fig.5 Network structure diagram of CBAM-YOLOv5

2.3 引入Transformer注意力機制網絡的改進型YOLOv5網絡設計

2.3.1 引入基礎型Transformer Encoder Block模塊的改進型YOLOv5網絡

Transformer編碼模塊(Transformer Encoder Block)模塊[10]采用原始Transformer網絡中的編碼(Encoder)結構，該Encoder網絡結構示意圖如圖6所示，它可以捕獲全局信息和豐富的上下文信息。

圖6 Transformer Encoder網絡結構示意圖 Fig.6 Network structure diagram of Transformer Encoder

每個Transformer Encoder塊包含兩個子層。第一子層為MHA層，第二子層為MLP層。每個子層之間使用殘差連接。

為了不大幅增加原始YOLOv5網絡的模型參數量，本文將Transformer Encoder結構僅插入在骨干特征提取網絡的末尾處，利用自注意力機制來捕獲最后一層特征提取層的局部信息特征。引入了基礎型Transformer Encoder Block模塊的改進型YOLOv5網絡如圖7所示，命名為BasicTrans-YOLOv5網絡。

圖7 BasicTrans-YOLOv5網絡結構示意圖Fig.7 Network structure diagram of BasicTrans-YOLOv5

2.3.2 引入Swin Transformer Block模塊的改進型YOLOv5網絡

Swin Transformer Block模塊是Swin Transformer網絡[11]的核心部分。Swin Transformer網絡是2021 年提出的一種采用了局部自注意力增強機制的Transformer網絡，該網絡相較于卷積神經網絡具有更強的動態計算能力，建模能力更強，可自適應計算局部與全局像素關系，非常具有推廣使用的價值。Swin Transformer Block模塊的網絡結構如圖8所示。

圖8 Swin Transformer Block網絡結構示意圖Fig.8 Network structure diagram of Swin Transformer Block

Transformer Block整體架構的核心模塊是基于窗口的多頭自注意力層(W-MSA)和基于偏移窗口的多頭自注意力層(SW-MSA)。網絡通過將注意力計算限制在一個窗口中，一方面能引入卷積操作的局部性，另一方面能節省計算量，具有很好的性能。

本文設計將Swin Transformer Block結構插入骨干特征提取網絡和頸部特征融合中，利用高效的自注意力機制模塊充分挖掘特征表征的潛能。引入Swin Transformer Block模塊的改進型YOLOv5網絡如圖9所示，命名為SwinTrans-YOLOv5網絡。

圖9 SwinTrans-YOLOv5網絡結構示意圖Fig.9 Network structure diagram of SwinTrans-YOLOv5

3 實驗結果及分析(Experimental results and analysis)

3.1 實驗環境

本文實驗采用私有的變電站內復雜多尺度檢測場景數據集，該數據集共有4 個分類，分別為屏體壓板、火災、抽煙、人員摔倒。統計樣本庫中各分類的樣本數量，梳理樣本庫中的正負樣本、難易樣本。針對少分類、難識別的圖像樣本，采用復制圖像后的隨機數據增強方法擴充樣本得到分類相對均衡的訓練樣本庫，增強處理后的數據集共包括10,497 張圖像。

使用開源工具對訓練樣本庫中待檢測目標的矩形關鍵區域進行人工標注，標注后使用腳本轉換形成YOLOv5算法支持的標注文件。訓練樣本庫按4:1的比例隨機建立訓練樣本集和測試樣本集。

3.2 結果對比與分析

本文設計的引入了注意力機制的改進型YOLOv5網絡包括兩大類、四種網絡結構，與基準的YOLOv5網絡在同一實驗平臺上進行訓練，模型評估結果對比如表1所示，其中mAP@0.5表示交并比為0.5時的平均精度均值(mAP)指標，推理耗時表示模型處理一張圖片所需要的時間，運算量表示網絡模型的浮點運算量。

表1 改進型YOLOv5網絡性能指標對比Tab.1 Performance comparison of improved YOLOv5 networks

表1中加粗字體為當前網絡在所屬的注意力模型大類下準確率的最優值，可以看出，在當前實驗環境參數及復雜多尺度場景數據集下有如下實驗結論。

(1)引入了注意力機制的目標檢測網絡相比基準的YOLOv5網絡，其mAP指標值一般都有一定的提升，說明注意力機制網絡的確能有效提升模型的表達能力，但代價是增加了模型的復雜度。

(2)在引入的基于卷積神經網絡的注意力機制模塊中，對簡單分類(屏體壓板)提升效果有限，相較于基準網絡，AP最大可提升1.0%，而對復雜分類(火災、抽煙、人員摔倒)樣本中的AP指標值的提升更明顯。橫向對比來看，引入CA注意力機制模塊的網絡綜合表現最佳：它相比基準YOLOv5網絡模型雖然復雜度略有提升，單張圖像的推理耗時也增加了1.4 ms，但mAP指標值提升了4.1%，并且在復雜分類樣本中的AP值提升效果顯著(火災、抽煙、人員摔倒的平均準確率分別提升4.9%、2.9%、7.4%)。因此從整體來看，CA模塊最具引入價值。

(3)在引入的基于Transformer結構的自注意力模塊網絡中，SwinTrans-YOLOv5的mAP指標值顯著提升，整體mAP指標相比基準YOLOv5網絡提升了9.0%，尤其是在復雜分類下的AP值的提升最顯著(火災、抽煙、人員摔倒的平均準確率分別提升6.8%、10.5%、18.3%)，但代價是推理耗時及模型復雜度都增加較大，比較適合于算力充足、對實時性要求不高的場合。

(4)總體來說，在模型部署硬件算力足夠的情況下，基于Transformer結構的SwinTrans-YOLOv5相較于基于卷積神經網絡結構的CA-YOLOv5，擁有更強大的建模能力和更高的檢測精度，在網絡設計中更具引入價值；但是，在模型部署硬件算力一般、需要考慮推理的實時性的情況下，CAYOLOv5也是一種非常不錯的設計思路。

當然，以上實驗結論并不是絕對的，它是建立在當前實驗環境下的數據結果，不同的數據集，效果可能不同，但引入注意力機制模塊確實能為提升原有基準網絡的檢測性能提供一種可行的設計思路。注意力機制模塊能彌補卷積網絡局部性過強、全局性不足的問題，幫助獲取全局的上下文信息，具有讓模型看得更廣的能力，尤其在一些復雜多尺度場景下，對于難樣本分類的檢測準確率一般會取得較好的效果。

4 結論(Conclusion)

本文簡要闡述了基于卷積神經網絡和基于Transformer網絡的注意力機制模塊的網絡結構，提出了多種引入了注意力機制的改進型YOLOv5網絡的構建方法，在復雜多尺度目標檢測數據集下進行性能對比。結果表明，引入注意力機制模塊的改進型網絡相較于基準網絡，或多或少地獲得了準確率方面的增益效果。其中，引入CA注意力模塊和引入Swin Transformer自注意力模塊分別在基于卷積神經網絡和基于Transformer網絡類別下取得了最佳的性能提升，并且基于Swin Transformer的自注意力網絡在復雜場景下的建模能力優于傳統的基于傳統卷積神經網絡的注意力網絡。本文的研究成果為復雜場景下多尺度難樣本目標檢測網絡的建模設計提供了一種改進思路。