基于注意力機制的YOLOv5優(yōu)化模型

潘燁新,黃啟鵬,韋 超,楊 哲

(1.蘇州大學 計算機科學與技術學院,江蘇 蘇州 215006;2.省計算機信息處理技術重點實驗室,江蘇 蘇州 215006)

0 引 言

目標檢測是機器視覺領域重要的研究內容之一[1],目前主流的檢測模型分為單階段模型、雙階段模型以及基于Transformer解編碼結構的模型[2]。雙階段算法先提取候選區(qū)域再進行分類和回歸,如RCNN[3],Faster R-CNN[4]系列。這些方法在檢測精度上表現(xiàn)出色,但由于計算量較大,檢測速度較慢。單階段檢測算法無需提取候選區(qū)域, 2直接對每個特征圖進行回歸預測。經(jīng)典的單階段檢測算法有YOLO[5],SSD[6],FCOS[7]等系列算法,YOLO因檢測速度快被廣泛應用于工業(yè)和日常生活中。但由于YOLOv5使用的骨干網(wǎng)絡CSPDarknet-53[8]提取的特征圖尺寸較小、分辨率較低、像素感受野較大,導致小目標的定位性能較差,因此整體性能仍存在一定的優(yōu)化空間。同時整個網(wǎng)絡中主要負責提取圖像特征的是C3模塊,分布在網(wǎng)絡的骨干和頸部中。在骨干部分,C3模塊可以為特征圖提取到大量的位置與細節(jié)信息,但語義信息提取的較少。當特征圖前向傳播到頸部部分后,在特征金字塔網(wǎng)絡(FPN[9])與像素聚合網(wǎng)絡(PAN)框架的結合作用下,C3模塊主要負責紋理特征的提取,此時會獲得較為豐富的語義信息,但丟失了大量的位置與細節(jié)信息。導致網(wǎng)絡模型對于小物體及有遮擋目標產(chǎn)生漏檢誤檢情況,性能下降。該文提出一種基于注意力機制的YOLOv5優(yōu)化模型。通過引入DRA(Dimension Related Attention,維度關聯(lián)注意力)模塊來解決C3模塊信息丟失問題,增強主干網(wǎng)絡提取圖像特征的能力;針對感受野大而導致的定位困難問題,引入新的定位計算損失函數(shù),在提高邊界框的定位精確度的同時優(yōu)化模型的推理速度,間接提升模型的性能。在通用數(shù)據(jù)集上的實驗結果表明,該方法提升了主干網(wǎng)絡的特征提取能力,降低了回歸參數(shù)的損失,從而提升了模型的整體性能。

1 相關工作

1.1 目標檢測模型

雙階段模型的代表RCNN開創(chuàng)性地使用深度學習模型進行目標檢測。但存在兩個問題:一是經(jīng)過縮放處理后會使一些圖片特征信息丟失,從而降低檢測的準確性,不利于小目標的檢測;二是在訓練和預測中,RCNN的速度都非常慢。Faster R-CNN提出了區(qū)域生成網(wǎng)絡(Region Proposal Networks,RPN)用于提升檢測框的生成速度,最終精度較高,但實時性與檢測小目標的效果差。YOLO是單階段模型的起始作,不再生成候選區(qū)而是直接進行分類和回歸。v1通過將圖像劃分成多個網(wǎng)格來生成候選框。相比于二階段模型,檢測速度有了很大提高,但精度相對較低,尤其在小目標檢測方面。v2[10]改變了主干網(wǎng)絡,相比v1模型在精度、速度和分類數(shù)量上都有了很大的改進,但由于每個網(wǎng)格只能預測一個物體,當同一個網(wǎng)格內包含多個物體時只能檢測到一個,因此對小物體的識別效果仍然非常差。v3[11]中提出了基于錨框的思想,使得最后的特征圖上基于每個單元格都有三個不同的尺寸大小的錨框,進而對錨框進行分類與回歸。v4[8]針對預處理以及激活函數(shù)問題,分別引入了Mosaic數(shù)據(jù)增強手段以及Mish激活函數(shù)[12],使得網(wǎng)絡的收斂速度與精度進一步提升,但仍然存在框定位不準以及召回率低的問題。YOLOv5在對模型主干以及頸部的基礎改進之外,更換了新的損失函數(shù)計算方法,同時優(yōu)化了一直存在的正負樣本分配問題。但對于整體而言,預測框的回歸精度與速度仍然較差。研究者們針對不同應用場景和問題,提出了基于YOLOv5的一系列應用優(yōu)化算法。張浩等人[13]提出的算法旨在提高無人機視角下密集小目標的檢測精度,并保證實時性。李永軍等人[14]將紅外成像與v5模型相結合,解決動態(tài)識別與密集目標的問題。竇其龍[15]通過優(yōu)化深度學習網(wǎng)絡、重新設置錨點框大小和嵌入GDAL模塊,提高檢測速度和降低漏檢率。劉閃亮[16]則提出了注意力特征融合結構,進一步提高模型對小目標的檢測性能。田楓[17]提出了Cascade-YOLOv5,用于油田場景規(guī)范化著裝檢測,來提高檢測性能。這些算法都是基于YOLOv5的改進和優(yōu)化,以適應不同領域和應用需求。

1.2 注意力機制

在機器視覺領域,常使用的是軟注意力,對其按維度可劃分為通道注意力、空間注意力和自注意力。通道注意力旨在聯(lián)系不同特征圖,通過網(wǎng)絡訓練獲取每個通道的重要度從而賦予不同權重最終強化重要特征,代表模型如SE-Net(Squeeze and Excitation)[18]。空間注意力通過空間轉換和掩碼加權等方式增強興趣區(qū)域[19]的同時弱化背景區(qū)域。如輕量級注意力模塊CBAM[20]。自注意力旨在最大化利用特征自身的固有信息進行交互。在Google提出的Transformer架構中被實際應用,何凱明等人將其應用到CV領域,并提出了Non-Local模塊[21],通過自注意力機制有效地捕獲長距離的特征依賴,實現(xiàn)全局上下文信息的建模。注意力機制模塊眾多,模型性能差異大,對比評估一些新型且有效的注意力機制模塊,并進行一些創(chuàng)新改進,對提升復雜多尺度目標的檢測性能是非常有意義的。

2 改進后的YOLOv5優(yōu)化模型

2.1 DRA注意力機制

DRA模塊在經(jīng)典的SE模塊上做出優(yōu)化,如公式1所示,它可以對網(wǎng)絡中任意的中間特征張量進行轉化變換后輸出同樣尺寸的張量。DRA模塊結構如圖1所示。

圖1 DRA注意力機制

X=[x1,x2,…,xc]∈RH×W×C→

Y=[y1,y2,…,yc]∈RH×W×C

(1)

在原先同時關注空間和通道信息的基礎上,通過改變全局池化的操作,保留通道間信息的同時考慮重要的空間信息。

通道注意力常采用全局池化編碼全局空間信息,簡而言之是全局信息被壓縮成了一個標量,而壓縮完之后的標量難以保留重要的空間信息。為解決此問題,DRA將全局池化操作改進為兩個1維向量的編碼操作。

為了獲取輸入圖像的高度與寬度上的注意力,并完成對精確位置信息的編碼,對于輸入特征圖,使用池化核(1,W)和(H,1)分別對高度和寬度的特征進行平均池化,從而獲得兩個方向的特征圖,如式2和式3所示。

(2)

(3)

對比全局池化的壓縮方式,這樣能夠允許注意力模塊捕捉單方向上的長距離關系,同時保留另一個方向上的空間信息,幫助網(wǎng)絡模型更準確地定位目標。

接著將獲得全局感受野的高度和寬度兩個方向的特征圖按通道維度拼接在一起,主要目的是方便之后進行批量歸一化(Batch Normalization,BN)操作。將它們送入卷積核為1×1的共享卷積模塊Conv2D,將其維度降低為C/r,r為可設定的縮減因子,接著對其進行BN處理,將得到的特征圖記為F1,最后送入Swish激活函數(shù)進行非線性變換,將這種變換記為δ,即可得到尺寸為C/r×1×(W+H)的包含橫向和縱向空間信息的特征圖f,如公式4所示。

f=δ(F1([Zh,Zw]))

(4)

隨后將f按照原來的高度和寬度進行卷積核大小為1×1的卷積,分別得到通道數(shù)與原來一樣的兩個獨立的特征fh和fw,最后經(jīng)過Sigmoid激活函數(shù)后,分別得到特征圖在高度上的注意力權重gh和在寬度方向的注意力權重gw,如式5和式6所示。

gh=σ(Fh(fh))

(5)

gw=σ(Fw(fw))

(6)

最后在原始特征圖上通過乘法加權計算,得到最終在寬度和高度方向上帶有注意力權重的特征圖,如式7所示。

(7)

2.2 注意力機制融合

針對原模型對于特征表達能力的不足,不易識別難檢目標,以及由于只考慮通道信息而缺失方向相關信息帶來的定位不準等問題,通過將DRA模塊插入到網(wǎng)絡模塊中解決。同時由于原模型的主干,頸部,檢測頭三層結構會帶來結果的干擾性以及不確定性,該文分析了在三層結構不同位置插入DRA模塊的效果。如圖2所示,分別在主干,頸部,預測頭中插入DRA模塊。對于主干部分,細分了DRA的插入位置。根據(jù)后續(xù)實驗表1的數(shù)據(jù),最終確定選擇(b)方式插入到主干,將新的整體結構命名為CDRA模塊,取代原模型主干中的C3模塊。CDRA模塊相比YOLO v5原先的C3模塊最大的改進在于,每個權重都包含了通道間信息、橫向空間信息和縱向空間信息,能夠幫助主干網(wǎng)絡更準確地定位目標信息,增強識別能力。

表1 不同位置的注意力機制融合結果對比 %

(a)原C3結構 (b)插入主干C3最后一層 (c)插入主干C3的殘差模塊 (d)插入Neck部分 (e)插入Head部分

將運用維度關聯(lián)注意力機制融合的CDRA替換原C3模塊,經(jīng)過改進后的YOLOv5s結構如圖3所示。圖3中,YOLOv5s主要由主干網(wǎng)絡、頸部、預測頭部網(wǎng)絡三部分組成,主干部分的替換工作對改進后的YOLOv5s性能提升起到?jīng)Q定性作用。

2.3 損失函數(shù)

目標檢測模型的損失函數(shù)通常由三個部分構成,分別為預測框的定位損失Lbox,置信度損失Lobj,分類損失Lcls,整體的網(wǎng)絡損失的計算如式8所示。

L=Lbox+Lobj+Lcls

(8)

其中置信度損失和分類損失均采用交叉熵損失(Binary Cross Entropy Loss),公式如式9所示。

(9)

預測框的定位損失用來衡量當前模型所給出的預測框與真實框之間位置上的誤差,具體會計算兩者的中心坐標、高、寬等誤差。早期模型一般采用L1,L2,smooth L1來計算該損失,但其忽略了4個回歸參數(shù)之間的相關性。當前常用的是交并比損失(Intersection over Union,IoU),IoU的計算公式如式10所示。

(10)

其中,B=(x,y,w,h)表示預測框的位置,Bgt=(xgt,ygt,wgt,hgt)表示真實框的位置。

IoU損失的計算公式如式11所示。

lossIoU=1-IoU

(11)

由公式可知兩個矩形框重合度越高,IoU越接近1,則損失越接近0。但采用IoU損失生效的情況僅在兩框之間有重疊的部分,對于非重疊的兩個框,IoU損失不會提供任何可供傳遞的梯度。

YOLOv5原始模型中采用CIoU作為邊界框的定位損失函數(shù)。CIoU是在DIoU(Distance IoU)[22]的基礎上考慮了兩框的長寬比而演化而來,但是仍然沒有考慮到真實框與預測框之間不匹配的方向。這種不足導致CIoU收斂速度較慢且效率較低。

為了解決CIoU存在的問題,該文引入SIoU[23]用以改進,保留了原損失函數(shù)的全部性質,同時考慮方向框的角度回歸問題,重新定義了懲罰指標。

SIoU由四部分組成:角度損失Λ、距離損失Δ、形狀損失Ω以及交并比損失(IoU)。

角度損失函數(shù)組件Λ,如式12所示。

(12)

其中,x是直角三角形中的對邊比斜邊,如圖4所示,α是兩框中心連線與預測框中心水平線的夾角。則x可由式13表示。

圖4 損失函數(shù)組件示意圖

(13)

其中,ch為真實框和預測框中心點的高度差,σ為真實框和預測框中心點的距離,可由式14表示。

(14)

將式13帶入式12化簡可得最終的角度損失計算公式,如式15。

Λ=sin(2α)

(15)

定義角度損失后,考慮到當出現(xiàn)同時存在一個角度很小但是很近,與一個角度很大但是很遠的框的情況時,近的框總是會被優(yōu)先選擇,所以直接使用角度損失不合理,還需要考慮距離與角度的互相關系。為了保證距離與角度的平衡,將角度損失同時考慮,重新定義了距離損失函數(shù),記為Δ,如式16所示。

微庫侖綜合分析儀（江蘇江分，JF-WK-2000）；輕油氯標準物質（10mg·kg-1，江蘇江分）；冰醋酸（優(yōu)級純，科密歐）；二次蒸餾水或去離子水。工作時參數(shù)設置。見表1。

(16)

其中,ρx,ρy,γ定義如式17所示。

(17)

ρt是使用原始的距離損失的平方次冪來賦權重,說明距離的影響要大于角度的影響。

形狀損失主要負責從長寬角度評價預測框的回歸參數(shù)與真實框是否相似,記為Ω,如式18所示。

(18)

其中,ωw,ωh如式19所示。

(19)

θ用來控制整體對形狀損失的關注程度。

綜合考慮上述3項以及默認的IoU損失,就可以得到最后的預測框的定位損失函數(shù),如式20所示。

(20)

3 實驗與分析

3.1 實驗環(huán)境

實驗環(huán)境配置:Window10操作系統(tǒng),32核Intel CPU,32 GB內存,兩塊TELSA A100,40 GB存儲空間。深度學習框架為PyTorch1.10,圖形處理器驅動為CUDA11.4和Cudnn8。訓練過程中所使用的優(yōu)化器為Adam[24],初始學習率為0.01,動量因子為0.937,權重衰減為0.000 25,批尺寸為32,總迭代次數(shù)設置為300。

3.2 數(shù)據(jù)集及預處理

使用Pascal VOC07+12訓練集以及VOC07測試數(shù)據(jù)集來評估模型性能,包含20個類別的常見交通工具、家具和動物等圖像,可用于目標檢測任務。共包含8 281張訓練圖像、8 333張驗證圖像和4 952張測試圖像。同時,在ImageNet數(shù)據(jù)集上對模型的主干網(wǎng)絡進行了預訓練,在訓練過程中,使用Mosaic數(shù)據(jù)增強技術對前75%的訓練周期進行了處理。

3.3 評估指標

使用檢測速度、檢測精度和損失函數(shù)收斂曲線等客觀指標來評價模型的性能。其中,FPS是檢測速度的評價指標。AP(Average Precision)是指在0～1范圍內P(Precision,正確率)指標對R(Recall,召回率)指標的積分,即P-R曲線下面積,AP值越大,模型精度越高。mAP是平均精度均值,指的是每個目標類別AP的均值。

(21)

(22)

(23)

(24)

式中,TP表示正確識別的目標數(shù)量,FP表示識別錯誤的目標數(shù)量,FN表示未被識別出目標數(shù)量。如果IoU大于一定閾值,則檢測框被標記為TP,否則為FP,如果檢測到真實目標沒有匹配到對應的預測框則標記為FN。

3.4 結果分析

3.4.1 改進模型的性能綜合分析

如2.3節(jié)所述,該文嘗試將DRA模塊融合到網(wǎng)絡模塊的不同位置,并對相應檢測結果展開對比。分別在原模型的主干,頸部,檢測頭中融入DRA模塊。特殊的對于backbone部分,更細化地對比了簡單的拼接在尾部或是融入原本的C3模塊中的結果數(shù)據(jù)。實驗結果如表1所示,將DRA模塊融入主干網(wǎng)絡中C3模塊的最后一層檢測效果最佳。YOLOv5網(wǎng)絡中提取特征的關鍵網(wǎng)絡在主干部分,其中隱含著易被網(wǎng)絡忽視掉的小目標特征信息,而在加入DRA模塊后,對這部分的特征信息進行了注意力重構,突出了重要信息,而在網(wǎng)絡更深的Neck以及Head部分,小目標的特征信息被淹沒,語義信息較為粗糙,注意力模塊難以區(qū)分出空間以及通道特征,自然無法很好地對特征進一步加強重構。

同時,將文中對YOLOv5的注意力及結合方式與其他注意力機制做對比,對比結果如表2所示,SE[18]是經(jīng)典的注意力機制起源,CA[25]是坐標注意力機制,CBAM[20]是經(jīng)典的空間通道注意力機制,ANG是一種輕量型的融合注意力機制方法模型?？梢钥闯瞿Ｐ筒⒉贿m合簡單地嵌套所有的注意力機制,當融合SE后,模型的漏檢率不降反增,說明網(wǎng)絡對于深層信息還是沒有掌握能力,再看ANG模型,輕量化的同時也帶來了精度的大量犧牲,而傳統(tǒng)的CA,CBAM也都基本維持在原精度附近,說明對于網(wǎng)絡沒有實質性的提升。

表2 不同注意力機制融合對比結果 %

為了分析不同的改進策略對于模型最后的檢測性能的影響,設計了4組消融實驗,結果如表3所示,其中,“×”代表在網(wǎng)絡中未使用的改進策略,“√”代表使用了改進。改進1在網(wǎng)絡中替換了損失函數(shù),解決了目標框與預測框的角度問題,使模型收斂速度與定位精準度提升;改進2在網(wǎng)絡主干部分的C3模塊中融合了DRA注意力機制,使得權重中同時包含了通道信息,橫向以及縱向空間信息,mAP提升了4.0百分點,FPS下降了12.8;改進3將兩者同時融入網(wǎng)絡中,如前文所述,模型在更好地提取特征的同時加快了收斂速度,mAP最終提升了4.7百分點,檢測速度則在改進2的基礎上加速了5.2,僅與原模型相差7.6。

表3 消融實驗結果

同時將消融實驗的mAP@0.5曲線繪制在同一個坐標系中,如圖5所示,改進后的模型在迭代次數(shù)達到45時逐漸趨于穩(wěn)定。進一步分析SIoU改進的數(shù)據(jù)值曲線,與原始模型的曲線對比,以更高的收斂速度趨于穩(wěn)定,表明了SIoU損失函數(shù)的替換使得回歸目標框能夠以更快的速度,更低的損失,精準地定位到待檢測目標。

圖5 不同改進策略的mAP@0.5對比

3.4.2 模型對比實驗

同時將文中模型與其他模型對所有類別檢測精度進行對比分析。Faster R-CNN[4]是二階段檢測模型的典型代表,SSD[6]是經(jīng)典單階段模型,v3[11]在精度和速度上有較好的均衡性能。v7[26]是當前表現(xiàn)較為出色的檢測模型,而YOLOv5是文中改進對象。表4為所比較模型在所有類別上的檢測平均精度對比,在所有20類上的檢測結果均優(yōu)于原v5s模型,平均精度均值為87.8%,同時與當前較為優(yōu)秀的YOLOv7模型相比,20類中有16類的結果高于v7,同時最終的平均精度均值提升1.4百分點。

表4 VOC上各類別平均精度結果

為進一步證實文中算法的有效性和優(yōu)越性,將文中算法模型與主流模型進行對比。由表5中實驗結果可得,文中算法模型在保持一定檢測速度的情況下,擁有更高的檢測精度。與傳統(tǒng)的雙階段算法Faster R-CNN[4]相比具有較大的檢測速度優(yōu)勢,平均精度均值提升了14.6百分點。與YOLO系列算法相比,改進模型比v3[11],v4[8],v5,v7[26]原始模型的檢測精度分別提高了10.6,15.1,4.7和1.7百分點。對于衡量難檢目標以及小目標檢測問題的閾值為0.5到0.95的平均精度均值(mAP@0.5:0.95),對比v5提升了4.5百分點,對比v7提升了2.8百分點。而在檢測速度方面,文中模型雖比原始模型有所降低,但仍達到83.3 frame/s,完全可以滿足工業(yè)場景下的實時檢測要求(30 frame/s)。

表5 不同模型的VOC數(shù)據(jù)集測試結果

4 結束語

YOLO系列目標檢測算法是運用較為廣泛的單階段目標檢測算法之一。針對YOLOv5對難檢目標,包括小目標和遮擋目標等檢測精度不高的問題,提出了注意力機制融合的方法,將DRA模塊與v5網(wǎng)絡的主干部分進行結合,以增強模型對于一些易漏信息的捕捉能力。同時使用了SIoU函數(shù)替換原損失函數(shù)中負責計算回歸參數(shù)的CIoU損失,提高了收斂速度和回歸精度,改善了遮擋等復雜情況下的漏檢以及小目標物體識別差的問題。實驗結果表明,改進模型的平均精度超越了原YOLOv5網(wǎng)絡。雖然模型參數(shù)量稍有增加,但改進模型的檢測速度仍符合工業(yè)需求的檢測速度。在后期研究中,還可以嘗試對于主干網(wǎng)絡中的卷積部分進行替換,或是替換特征加強的Neck部分,進一步提升模型對于難檢目標的檢測精度。