引入注意力機(jī)制的改進(jìn)型YOLOv5網(wǎng)絡(luò)研究

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

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

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

1 引言(Introduction)

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

注意力機(jī)制可分為基于卷積神經(jīng)網(wǎng)絡(luò)的注意力機(jī)制[6]和基于Transformer網(wǎng)絡(luò)的自注意力機(jī)制[7]兩大類。注意力機(jī)制模塊眾多、模型性能差異大，對(duì)比評(píng)估一些新型且有效的注意力機(jī)制模塊，對(duì)提升復(fù)雜多尺度目標(biāo)的檢測(cè)性能是非常有意義的。

本文選取YOLOv5網(wǎng)絡(luò)作為基準(zhǔn)模型，在網(wǎng)絡(luò)中引入多種先進(jìn)的注意力機(jī)制模塊，設(shè)計(jì)并實(shí)驗(yàn)對(duì)比分析了幾種改進(jìn)型網(wǎng)絡(luò)的性能表現(xiàn)。

2 改進(jìn)型YOLOv5網(wǎng)絡(luò)的設(shè)計(jì)(Design of improved YOLOv5 network)

2.1 標(biāo)準(zhǔn)YOLOv5目標(biāo)檢測(cè)網(wǎng)絡(luò)結(jié)構(gòu)

YOLOv5網(wǎng)絡(luò)在工業(yè)落地應(yīng)用中表現(xiàn)出極其優(yōu)秀的檢測(cè)性能和推廣價(jià)值，它整合了大量的計(jì)算機(jī)視覺(jué)前沿技術(shù)，顯著改善了對(duì)象檢測(cè)的性能，提升了模型訓(xùn)練的速度及模型應(yīng)用的便利度。

YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)示意圖如圖1所示。

圖1 YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)示意圖Fig.1 Diagram of YOLOv5 network structure

圖1中的基準(zhǔn)網(wǎng)絡(luò)主要由骨干特征提取網(wǎng)絡(luò)(Backbone網(wǎng)絡(luò))、頸部特征融合網(wǎng)絡(luò)(Neck網(wǎng)絡(luò))和檢測(cè)頭部預(yù)測(cè)網(wǎng)絡(luò)(Prediction網(wǎng)絡(luò))組成，分別主要使用基于瓶頸(Bottleneck)結(jié)構(gòu)的C3為Backbone，基于多特征圖融合的FPN+PANet結(jié)構(gòu)為Neck，以及基于檢測(cè)目標(biāo)的位置和類別進(jìn)行回歸、分類任務(wù)的YOLO檢測(cè)頭為Prediction。

2.2 引入卷積神經(jīng)網(wǎng)絡(luò)注意力機(jī)制的改進(jìn)型YOLOv5網(wǎng)絡(luò)設(shè)計(jì)

2.2.1 引入?yún)f(xié)同注意力機(jī)制網(wǎng)絡(luò)的改進(jìn)型YOLOv5網(wǎng)絡(luò)

協(xié)同注意力機(jī)制網(wǎng)絡(luò)(Coordinate Attention，CA)[8]，它引入了一種新的注意塊結(jié)構(gòu)，該結(jié)構(gòu)不僅能捕獲跨通道的信息，還能捕獲方向感知和位置感知的信息，這能幫助模型更加精準(zhǔn)地定位和識(shí)別感興趣的目標(biāo)。

基于CA模塊的典型網(wǎng)絡(luò)結(jié)構(gòu)如圖2所示。

圖2 CA注意力模塊的典型網(wǎng)絡(luò)結(jié)構(gòu)示意圖Fig.2 Network structure diagram of CA attention module

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

本文將CA網(wǎng)絡(luò)模塊應(yīng)用在YOLOv5網(wǎng)絡(luò)的骨干特征提取網(wǎng)絡(luò)的每個(gè)多尺度特征輸出的位置，用于對(duì)每個(gè)尺度下特征圖的各通道進(jìn)行特征重標(biāo)定，以提升原始YOLOv5網(wǎng)絡(luò)的特征提取能力。設(shè)計(jì)的引入CA網(wǎng)絡(luò)模塊的改進(jìn)型YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)示意圖如圖3所示，命名為CA-YOLOv5網(wǎng)絡(luò)。

圖3 CA-YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)示意圖Fig.3 Diagram of CA-YOLOv5 network structure

2.2.2 引入CBAM網(wǎng)絡(luò)的改進(jìn)型YOLOv5網(wǎng)絡(luò)

CBAM網(wǎng)絡(luò)在通道注意力網(wǎng)絡(luò)的基礎(chǔ)上擴(kuò)展了一個(gè)空間注意力模塊，它可以在通道和空間維度上進(jìn)行注意力運(yùn)算。CBAM網(wǎng)絡(luò)[9]包含兩個(gè)子模塊：通道注意力模塊(Channel Attention Module，CAM)和空間注意力模塊(Spatial Attention Module，SAM)。

基于CBAM注意力模塊的典型網(wǎng)絡(luò)示意圖如圖4所示。

圖4 基于CBAM注意力模塊的典型網(wǎng)絡(luò)示意圖Fig.4 Network structure diagram of CBAM attention module

CAM模塊通過(guò)一個(gè)并行的最大值池化層，得到兩個(gè)1×1×C的特征圖，首先將得到的兩個(gè)特征向量相加，然后經(jīng)過(guò)一個(gè)激活函數(shù)得到每個(gè)通道的權(quán)重系數(shù)，最后用權(quán)重系數(shù)與原來(lái)的特征圖通道相乘，即可得到縮放后的新特征圖。

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

本文將CBAM網(wǎng)絡(luò)模塊應(yīng)用在YOLOv5網(wǎng)絡(luò)的骨干特征提取網(wǎng)絡(luò)的每個(gè)多尺度特征輸出的位置，用于對(duì)每個(gè)尺度下特征圖的各通道/各空間進(jìn)行特征重標(biāo)定，以提升原始YOLOv5網(wǎng)絡(luò)的特征提取能力。設(shè)計(jì)引入CBAM網(wǎng)絡(luò)的改進(jìn)型YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)示意圖如圖5所示，命名為CBAMYOLOv5網(wǎng)絡(luò)。

圖5 CBAM-YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)示意圖Fig.5 Network structure diagram of CBAM-YOLOv5

2.3 引入Transformer注意力機(jī)制網(wǎng)絡(luò)的改進(jìn)型YOLOv5網(wǎng)絡(luò)設(shè)計(jì)

2.3.1 引入基礎(chǔ)型Transformer Encoder Block模塊的改進(jìn)型YOLOv5網(wǎng)絡(luò)

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

圖6 Transformer Encoder網(wǎng)絡(luò)結(jié)構(gòu)示意圖 Fig.6 Network structure diagram of Transformer Encoder

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

為了不大幅增加原始YOLOv5網(wǎng)絡(luò)的模型參數(shù)量，本文將Transformer Encoder結(jié)構(gòu)僅插入在骨干特征提取網(wǎng)絡(luò)的末尾處，利用自注意力機(jī)制來(lái)捕獲最后一層特征提取層的局部信息特征。引入了基礎(chǔ)型Transformer Encoder Block模塊的改進(jìn)型YOLOv5網(wǎng)絡(luò)如圖7所示，命名為BasicTrans-YOLOv5網(wǎng)絡(luò)。

圖7 BasicTrans-YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)示意圖Fig.7 Network structure diagram of BasicTrans-YOLOv5

2.3.2 引入Swin Transformer Block模塊的改進(jìn)型YOLOv5網(wǎng)絡(luò)

Swin Transformer Block模塊是Swin Transformer網(wǎng)絡(luò)[11]的核心部分。Swin Transformer網(wǎng)絡(luò)是2021 年提出的一種采用了局部自注意力增強(qiáng)機(jī)制的Transformer網(wǎng)絡(luò)，該網(wǎng)絡(luò)相較于卷積神經(jīng)網(wǎng)絡(luò)具有更強(qiáng)的動(dòng)態(tài)計(jì)算能力，建模能力更強(qiáng)，可自適應(yīng)計(jì)算局部與全局像素關(guān)系，非常具有推廣使用的價(jià)值。Swin Transformer Block模塊的網(wǎng)絡(luò)結(jié)構(gòu)如圖8所示。

圖8 Swin Transformer Block網(wǎng)絡(luò)結(jié)構(gòu)示意圖Fig.8 Network structure diagram of Swin Transformer Block

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

本文設(shè)計(jì)將Swin Transformer Block結(jié)構(gòu)插入骨干特征提取網(wǎng)絡(luò)和頸部特征融合中，利用高效的自注意力機(jī)制模塊充分挖掘特征表征的潛能。引入Swin Transformer Block模塊的改進(jìn)型YOLOv5網(wǎng)絡(luò)如圖9所示，命名為SwinTrans-YOLOv5網(wǎng)絡(luò)。

圖9 SwinTrans-YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)示意圖Fig.9 Network structure diagram of SwinTrans-YOLOv5

3 實(shí)驗(yàn)結(jié)果及分析(Experimental results and analysis)

3.1 實(shí)驗(yàn)環(huán)境

本文實(shí)驗(yàn)采用私有的變電站內(nèi)復(fù)雜多尺度檢測(cè)場(chǎng)景數(shù)據(jù)集，該數(shù)據(jù)集共有4 個(gè)分類，分別為屏體壓板、火災(zāi)、抽煙、人員摔倒。統(tǒng)計(jì)樣本庫(kù)中各分類的樣本數(shù)量，梳理樣本庫(kù)中的正負(fù)樣本、難易樣本。針對(duì)少分類、難識(shí)別的圖像樣本，采用復(fù)制圖像后的隨機(jī)數(shù)據(jù)增強(qiáng)方法擴(kuò)充樣本得到分類相對(duì)均衡的訓(xùn)練樣本庫(kù)，增強(qiáng)處理后的數(shù)據(jù)集共包括10,497 張圖像。

使用開(kāi)源工具對(duì)訓(xùn)練樣本庫(kù)中待檢測(cè)目標(biāo)的矩形關(guān)鍵區(qū)域進(jìn)行人工標(biāo)注，標(biāo)注后使用腳本轉(zhuǎn)換形成YOLOv5算法支持的標(biāo)注文件。訓(xùn)練樣本庫(kù)按4:1的比例隨機(jī)建立訓(xùn)練樣本集和測(cè)試樣本集。

3.2 結(jié)果對(duì)比與分析

本文設(shè)計(jì)的引入了注意力機(jī)制的改進(jìn)型YOLOv5網(wǎng)絡(luò)包括兩大類、四種網(wǎng)絡(luò)結(jié)構(gòu)，與基準(zhǔn)的YOLOv5網(wǎng)絡(luò)在同一實(shí)驗(yàn)平臺(tái)上進(jìn)行訓(xùn)練，模型評(píng)估結(jié)果對(duì)比如表1所示，其中mAP@0.5表示交并比為0.5時(shí)的平均精度均值(mAP)指標(biāo)，推理耗時(shí)表示模型處理一張圖片所需要的時(shí)間，運(yùn)算量表示網(wǎng)絡(luò)模型的浮點(diǎn)運(yùn)算量。

表1 改進(jìn)型YOLOv5網(wǎng)絡(luò)性能指標(biāo)對(duì)比Tab.1 Performance comparison of improved YOLOv5 networks

表1中加粗字體為當(dāng)前網(wǎng)絡(luò)在所屬的注意力模型大類下準(zhǔn)確率的最優(yōu)值，可以看出，在當(dāng)前實(shí)驗(yàn)環(huán)境參數(shù)及復(fù)雜多尺度場(chǎng)景數(shù)據(jù)集下有如下實(shí)驗(yàn)結(jié)論。

(1)引入了注意力機(jī)制的目標(biāo)檢測(cè)網(wǎng)絡(luò)相比基準(zhǔn)的YOLOv5網(wǎng)絡(luò)，其mAP指標(biāo)值一般都有一定的提升，說(shuō)明注意力機(jī)制網(wǎng)絡(luò)的確能有效提升模型的表達(dá)能力，但代價(jià)是增加了模型的復(fù)雜度。

(2)在引入的基于卷積神經(jīng)網(wǎng)絡(luò)的注意力機(jī)制模塊中，對(duì)簡(jiǎn)單分類(屏體壓板)提升效果有限，相較于基準(zhǔn)網(wǎng)絡(luò)，AP最大可提升1.0%，而對(duì)復(fù)雜分類(火災(zāi)、抽煙、人員摔倒)樣本中的AP指標(biāo)值的提升更明顯。橫向?qū)Ρ葋?lái)看，引入CA注意力機(jī)制模塊的網(wǎng)絡(luò)綜合表現(xiàn)最佳：它相比基準(zhǔn)YOLOv5網(wǎng)絡(luò)模型雖然復(fù)雜度略有提升，單張圖像的推理耗時(shí)也增加了1.4 ms，但mAP指標(biāo)值提升了4.1%，并且在復(fù)雜分類樣本中的AP值提升效果顯著(火災(zāi)、抽煙、人員摔倒的平均準(zhǔn)確率分別提升4.9%、2.9%、7.4%)。因此從整體來(lái)看，CA模塊最具引入價(jià)值。

(3)在引入的基于Transformer結(jié)構(gòu)的自注意力模塊網(wǎng)絡(luò)中，SwinTrans-YOLOv5的mAP指標(biāo)值顯著提升，整體mAP指標(biāo)相比基準(zhǔn)YOLOv5網(wǎng)絡(luò)提升了9.0%，尤其是在復(fù)雜分類下的AP值的提升最顯著(火災(zāi)、抽煙、人員摔倒的平均準(zhǔn)確率分別提升6.8%、10.5%、18.3%)，但代價(jià)是推理耗時(shí)及模型復(fù)雜度都增加較大，比較適合于算力充足、對(duì)實(shí)時(shí)性要求不高的場(chǎng)合。

(4)總體來(lái)說(shuō)，在模型部署硬件算力足夠的情況下，基于Transformer結(jié)構(gòu)的SwinTrans-YOLOv5相較于基于卷積神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)的CA-YOLOv5，擁有更強(qiáng)大的建模能力和更高的檢測(cè)精度，在網(wǎng)絡(luò)設(shè)計(jì)中更具引入價(jià)值；但是，在模型部署硬件算力一般、需要考慮推理的實(shí)時(shí)性的情況下，CAYOLOv5也是一種非常不錯(cuò)的設(shè)計(jì)思路。

當(dāng)然，以上實(shí)驗(yàn)結(jié)論并不是絕對(duì)的，它是建立在當(dāng)前實(shí)驗(yàn)環(huán)境下的數(shù)據(jù)結(jié)果，不同的數(shù)據(jù)集，效果可能不同，但引入注意力機(jī)制模塊確實(shí)能為提升原有基準(zhǔn)網(wǎng)絡(luò)的檢測(cè)性能提供一種可行的設(shè)計(jì)思路。注意力機(jī)制模塊能彌補(bǔ)卷積網(wǎng)絡(luò)局部性過(guò)強(qiáng)、全局性不足的問(wèn)題，幫助獲取全局的上下文信息，具有讓模型看得更廣的能力，尤其在一些復(fù)雜多尺度場(chǎng)景下，對(duì)于難樣本分類的檢測(cè)準(zhǔn)確率一般會(huì)取得較好的效果。

4 結(jié)論(Conclusion)

本文簡(jiǎn)要闡述了基于卷積神經(jīng)網(wǎng)絡(luò)和基于Transformer網(wǎng)絡(luò)的注意力機(jī)制模塊的網(wǎng)絡(luò)結(jié)構(gòu)，提出了多種引入了注意力機(jī)制的改進(jìn)型YOLOv5網(wǎng)絡(luò)的構(gòu)建方法，在復(fù)雜多尺度目標(biāo)檢測(cè)數(shù)據(jù)集下進(jìn)行性能對(duì)比。結(jié)果表明，引入注意力機(jī)制模塊的改進(jìn)型網(wǎng)絡(luò)相較于基準(zhǔn)網(wǎng)絡(luò)，或多或少地獲得了準(zhǔn)確率方面的增益效果。其中，引入CA注意力模塊和引入Swin Transformer自注意力模塊分別在基于卷積神經(jīng)網(wǎng)絡(luò)和基于Transformer網(wǎng)絡(luò)類別下取得了最佳的性能提升，并且基于Swin Transformer的自注意力網(wǎng)絡(luò)在復(fù)雜場(chǎng)景下的建模能力優(yōu)于傳統(tǒng)的基于傳統(tǒng)卷積神經(jīng)網(wǎng)絡(luò)的注意力網(wǎng)絡(luò)。本文的研究成果為復(fù)雜場(chǎng)景下多尺度難樣本目標(biāo)檢測(cè)網(wǎng)絡(luò)的建模設(shè)計(jì)提供了一種改進(jìn)思路。