融合CBAM的YOLOv4輕量化檢測方法

任豐儀,裴信彪,喬 正,白 越

1(中國科學(xué)院 長春光學(xué)精密機(jī)械與物理研究所,長春 130033) 2(中國科學(xué)院大學(xué),北京 100039)

1 引 言

伴隨著無人駕駛飛機(jī)的廣泛普及,促進(jìn)了城市管理[1],軍事偵察、災(zāi)害救援,土地變化監(jiān)控[2]和交通監(jiān)控[3]等多種應(yīng)用.過去10年見證了無人機(jī)視覺技術(shù)的巨大進(jìn)步,尤其是基于深度學(xué)習(xí)的目標(biāo)檢測算法,應(yīng)用于無人機(jī)視覺領(lǐng)域會極大提高無人機(jī)的場景理解能力,對于無人機(jī)獲取的航拍圖像進(jìn)行實(shí)時(shí)目標(biāo)檢測逐漸成為研究熱點(diǎn).

目標(biāo)檢測模型通常由3個部分組成:主干網(wǎng)絡(luò),頸部網(wǎng)絡(luò)和頭部網(wǎng)絡(luò).主干網(wǎng)絡(luò)功能是進(jìn)行初步特征提取,對于在GPU平臺上運(yùn)行的檢測模型,可以選用復(fù)雜度高的網(wǎng)絡(luò),如VGG,CSPDarknet53[4],DenseNet,ResNet或ResNeXt.對于在CPU平臺上運(yùn)行的檢測模型,要選擇緊湊網(wǎng)絡(luò),如SqueezeNet[5],MobileNet[6-8]或ShuffleNet[9].頸部網(wǎng)絡(luò)的作用是加強(qiáng)特征提取,具體包括FPN,PANet,Bi-FPN等模塊.頭部網(wǎng)絡(luò)利用得到的特征進(jìn)行預(yù)測,主要分為兩類,一類是基于區(qū)域預(yù)測的Two stage算法,另一類是基于回歸問題的One stage算法,Two stage算法先生成候選框,再利用CNN進(jìn)行特征提取與分類,其主要代表為R-CNN[10]系列,包括Fast R-CNN[11],Faster R-CNN[12],R-FCN[13]和Mask R-CNN[14].One stage算法不再采用候選框,而是直接對目標(biāo)物體邊界框及類別進(jìn)行回歸,代表算法為YOLO[15],SSD[16]和RetinaNet[17].總體上,前者相對精度更高,后者檢測速度更快.

近些年來,很多研究關(guān)注于在目標(biāo)檢測算法中添加功能模塊,從而在增加少量推理成本的同時(shí),提高目標(biāo)精度.增強(qiáng)感受野的常見模塊是SPP[18]、ASPP[19]和RFB[20];物體檢測中經(jīng)常使用的注意力模塊有SE[21]、ECA[22]和CBAM[23];用于篩選模型預(yù)測結(jié)果的常見后處理方法是NMS[24],但原始的NMS沒有考慮上下文信息,因此2017年提出了更加優(yōu)化的Soft-NMS[25]策略.

到目前為止,在無人機(jī)獲取的圖像上進(jìn)行實(shí)時(shí)目標(biāo)檢測面臨著挑戰(zhàn)和困難,首要就是深度神經(jīng)網(wǎng)絡(luò)存在復(fù)雜性和存儲量高的問題.本文針對這個問題,所做的工作主要分為3個方面:

1)以YOLOv4模型為基礎(chǔ),使用MobileNet模型替換YOLOv4本身的主干網(wǎng)絡(luò)CSPDarknet53,并利用MobileNet提出的深度可分離卷積思想,將原網(wǎng)絡(luò)中PANet與Head模塊的3×3標(biāo)準(zhǔn)卷積塊替換為深度可分離卷積塊.

2)在MobileNet-YOLOv4基礎(chǔ)上進(jìn)行模型優(yōu)化,加入CBAM注意力網(wǎng)絡(luò),并在算法后處理部分引入Soft-NMS模塊替代網(wǎng)絡(luò)原先的NMS模塊.在不影響模型運(yùn)行速度的前提下,提高模型的檢測精度.

3)將本文模型進(jìn)行訓(xùn)練后部署到無人機(jī)裝載的Nvidia Jetson TX2和Raspberry Pi低功耗嵌入式平臺,通過飛行試驗(yàn)實(shí)時(shí)定位和識別無人機(jī)航拍圖像中的車輛和行人.

2 改進(jìn)的輕量化YOLOv4模型

2.1 YOLOv4基本模型

如圖1所示,YOLOv4整體結(jié)構(gòu)可以拆分成3部分:

圖1 YOLOv4的整體結(jié)構(gòu)Fig.1 Overall structure of YOLOv4

1)主干網(wǎng)絡(luò):主干特征提取網(wǎng)絡(luò)選用CSPDarknet53,進(jìn)行初步特征提取.可以獲得3個初步的有效特征層,分別位于主干網(wǎng)絡(luò)的中間層、中下層、底層,3個特征層的大小分別為(52,52,256)、(26,26,512)、(13,13,1024),使用3個尺度的特征層進(jìn)行分類與回歸預(yù)測.

2)頸部網(wǎng)絡(luò):為了加強(qiáng)特征提取,從特征獲取預(yù)測結(jié)果的過程可以分為兩個部分,首先構(gòu)建SPP模塊、FPN+PAN特征金字塔結(jié)構(gòu)進(jìn)行加強(qiáng)特征提取;接下來利用YOLO Head對3個有效特征層進(jìn)行預(yù)測.YOLO Head本質(zhì)上是一次3×3卷積加上一次1×1卷積,作用分別是特征整合和調(diào)整通道數(shù),可以對3個初步的有效特征層進(jìn)行特征融合.

3)預(yù)測:第3部分為預(yù)測網(wǎng)絡(luò),利用更有效的特征層獲得預(yù)測結(jié)果.

2.2 MobileNet基本模型

2017年Google提出了MobileNet v1網(wǎng)絡(luò)結(jié)構(gòu),它使用了深度可分離卷積以及縮放因子,主要特點(diǎn)是模型小、計(jì)算速度快.MobileNet卷積神經(jīng)網(wǎng)絡(luò)極低的參數(shù)量和運(yùn)算量的優(yōu)點(diǎn)使其更適合部署在運(yùn)算量低的嵌入式設(shè)備上,可以更好的平衡檢測模型的準(zhǔn)確度和運(yùn)行的效率.

MobileNet v1是一種流水型網(wǎng)絡(luò)結(jié)構(gòu),它的主要特點(diǎn)分為兩方面:1)使用深度可分離卷積替代了傳統(tǒng)的卷積操作,構(gòu)建輕量級神經(jīng)網(wǎng)絡(luò).如圖2所示,深度級可分離卷積可以分解為兩個更小的操作:深度卷積和逐點(diǎn)卷積.深度卷積將卷積核拆分成單通道形式,對每一通道進(jìn)行卷積操作,得到與輸入特征圖通道數(shù)一致的輸出特征圖;逐點(diǎn)卷積就是1×1卷積,主要作用就是對特征圖進(jìn)行升維和降維;2)引入寬度α和分辨率ρ縮放因子.α對網(wǎng)絡(luò)輸入和輸出通道數(shù)進(jìn)行縮減,ρ用于控制輸入和內(nèi)部層表示,即控制輸入的分辨率,都可以進(jìn)一步縮小模型.

圖2 左圖:標(biāo)準(zhǔn)卷積,右圖:深度可分離卷積Fig.2 Left standard convolution rightdepth separable convolution

除此之外,加入了更多的ReLU6激活函數(shù),增加了模型的非線性變化,從而提高泛化能力.ReLU6函數(shù)與其導(dǎo)函數(shù)如下:

relu6(x)=min(max(x,0),6∈[0,6]

(1)

(2)

表1 深度可分離卷積與標(biāo)準(zhǔn)卷積的對比Table 1 Comparison of factorized convolutions and standard convolution

本文將MobileNet作為YOLOv4的主干特征提取網(wǎng)絡(luò),利用MobileNet模型強(qiáng)大特征提取能力、極低參數(shù)量和運(yùn)算量的優(yōu)勢,可以提高模型的運(yùn)算效率,并且將YOLOv4中的部分3×3標(biāo)準(zhǔn)卷積替換為深度可分離卷積,在降低模型計(jì)算量的同時(shí)提高性能.

2.3 MobileNet-YOLOv4優(yōu)化模型

近年來,有很多研究嘗試將注意力機(jī)制引入到卷積神經(jīng)網(wǎng)絡(luò)中,以提高其在大規(guī)模分類任務(wù)中的性能.CBAM是一種能對特征圖像局部信息聚焦的模塊.它通過學(xué)習(xí)的方式在空間和通道上對特征圖像進(jìn)行權(quán)重分配,促使計(jì)算資源更傾向于重點(diǎn)關(guān)注的目標(biāo)區(qū)域,從而加強(qiáng)感興趣的信息,同時(shí)抑制無用信息.CBAM 包含兩個模塊,輸入特征依次通過通道注意力模塊、空間注意力模塊的篩選,最后獲得經(jīng)過了重標(biāo)定的特征,即強(qiáng)調(diào)重要特征,壓縮不重要特征.模塊劃分如圖3所示.

圖3 CBAM模塊劃分Fig.3 CBAM module division

將 CBAM、SE注意力機(jī)制通過Grad-CAM[26]方法進(jìn)行可視化,并與MobileNet網(wǎng)絡(luò)的可視化結(jié)果進(jìn)行比較,結(jié)果如圖4所示.可以看出MobileNet與CBAM方法同時(shí)存在時(shí),Grad-CAM掩碼很好地覆蓋了目標(biāo)對象的區(qū)域,與SE算法相比,有效預(yù)測區(qū)域范圍更大,結(jié)果也更加準(zhǔn)確.使用CBAM注意力機(jī)制可以很好地學(xué)習(xí)利用目標(biāo)區(qū)域中的信息并從中聚合特征.

圖4 Grad-CAM可視化結(jié)果Fig.4 Grad-CAM visualization results

目標(biāo)檢測算法中,非極大值抑制策略(NMS)是很重要部分,對重疊框的處理方式如式(3)所示:

(3)

其中IoU表示重疊度,NMS保留在其閾值內(nèi)的檢測框,它的問題在于它會將與目標(biāo)框相鄰的檢測框的分?jǐn)?shù)強(qiáng)制歸零,導(dǎo)致漏檢和目標(biāo)定位錯誤,如圖5所示,對于置信度不高的目標(biāo)檢測結(jié)果,容易出現(xiàn)因圖像微小變化導(dǎo)致兩個預(yù)測框置信度大小關(guān)系產(chǎn)生變化,使得在NMS階段舍棄保留關(guān)系改變,最終第1 幀檢測結(jié)果偏左上方,第2幀偏右下方.

圖5 NMS損害預(yù)測框定位穩(wěn)定性的原因Fig.5 Reason for origimal NMS reducing the stability of bounding box

本文在YOLOv4模型中引入Soft-NMS算法來代替 NMS 算法.Soft-NMS同時(shí)考慮了得分和重合程度,對于與最高得分的檢測框重疊度較高的框設(shè)置一個懲罰項(xiàng),避免重疊框如果包含目標(biāo)卻被刪除造成漏檢的情況,同時(shí)不保留同一個目標(biāo)兩個相似的檢測框.處理方式如式(4)所示:

(4)

Soft-NMS實(shí)現(xiàn)過程如下:

Input:B={b1,…,bN},S={s1,…,sN},Nt

begin:

D←{}

whileB≠emptydo

m←argmaxS

M←bm

D←D∪M;B←B←M

forbiin Bdo

ifiou(M,bi)Ntthen

B←B-bi;S←S-si

endNMS

si←sif(iou(M,bi))

Soft-NMS

end

end

returnD,S

end

提出的模型以MobileNet-YOLOv4為特征提取網(wǎng)絡(luò),并且將PANet和YOLOHead中的標(biāo)準(zhǔn)卷積替換為深度可分離卷積,進(jìn)一步地減少計(jì)算量.此外,在模型3個不同尺度的預(yù)測部分加入CBAM注意力機(jī)制,從而能夠增強(qiáng)空間維度和通道維度的有效特征,抑制無效信息的流動,并且將非極大抑制階段的NMS 替換為Soft-NMS.模型整體框架如圖6所示.

圖6 輕量級網(wǎng)絡(luò)模型總體結(jié)構(gòu)Fig.6 Overall structure of lightweight network model

3 實(shí) 驗(yàn)

為了訓(xùn)練和評估所提出的輕量級模型,實(shí)驗(yàn)使用深度學(xué)習(xí)框架Pytorch,CPU選用Core i9-10900K,GPU選用Nvidia Geforce GTX 3080.此外,為了驗(yàn)證模型在無人機(jī)飛行時(shí)的適用性,在嵌入式系統(tǒng)Nvidia Jetson TX2和Raspberry Pi 4B上也進(jìn)行了模型部署和實(shí)驗(yàn)結(jié)果的分析.

3.1 實(shí)驗(yàn)訓(xùn)練數(shù)據(jù)集

一般模型在進(jìn)行訓(xùn)練和性能評估時(shí),有很多數(shù)據(jù)集可以選擇,其中最常用的是ImageNet數(shù)據(jù)集[27]、MS COCO數(shù)據(jù)集[28]和PASCAL VOC數(shù)據(jù)集[29].本次實(shí)驗(yàn)選用包含20個類別的PASCAL VOC作為模型訓(xùn)練和測試的數(shù)據(jù)集,劃分驗(yàn)證子集和訓(xùn)練子集的比例為1:9.為了降低各方面額外因素對識別的影響,對原始數(shù)據(jù)集進(jìn)行數(shù)據(jù)增強(qiáng).對構(gòu)建好的模型進(jìn)行訓(xùn)練微調(diào)時(shí),設(shè)置momentum=0.9,lr=0.001,batch_size=16,Init_Epoch=0,Freeze_Epoch=50,去除掉了優(yōu)化器的權(quán)重衰減因子,即weight decay=0.

3.2 實(shí)驗(yàn)過程與結(jié)果分析

通道剪枝和緊湊型網(wǎng)絡(luò)設(shè)計(jì),是輕量化網(wǎng)絡(luò)的常見方法.

圖7 YOLOv4通道剪枝流程圖Fig.7 Flow chart of YOLOv4 channel pruning

為了完成對比實(shí)驗(yàn),本文首先按照如圖7所示的步驟在 YOLOv4 中應(yīng)用通道修剪以獲取Slim YOLOv4.Network slimming[30]基本原理是對各個通道加入縮放因子,將其與通道的輸出相乘,通過訓(xùn)練進(jìn)行網(wǎng)絡(luò)權(quán)重和縮放因子的更新迭代,將小縮放因子對應(yīng)的通道進(jìn)行刪減,然后進(jìn)行網(wǎng)絡(luò)微調(diào),實(shí)現(xiàn)模型壓縮.緊湊型網(wǎng)絡(luò)實(shí)驗(yàn)是將MobileNet系列的3個網(wǎng)絡(luò)分別替換原始YOLOv4本身的主干網(wǎng)絡(luò)CSPDarkNet53,實(shí)驗(yàn)結(jié)果如表2所示.可以看出利用MobileNet模型強(qiáng)大特征提取能力、極低參數(shù)量和運(yùn)算量的優(yōu)勢,模型可以在小幅減少精度前提下,極大提高運(yùn)算效率,比通道剪枝的效果更優(yōu),并且MobileNetv1作為主干網(wǎng)絡(luò)時(shí)效果最佳.

表2 YOLOv4網(wǎng)絡(luò)輕量化前后對比Table 2 YOLOv4 network lightweight before and after comparison

接下來對MobileNet-YOLOv4模型進(jìn)行優(yōu)化.首先將模型PANet和YOLO Head部分中的3×3標(biāo)準(zhǔn)卷積替換為深度可分離卷積塊,再將CBAM卷積注意力機(jī)制嵌入MobileNet-YOLOv4模型,并在模型后處理階段引入 Soft-NMS 算法來提高網(wǎng)絡(luò)性能,NMS和Soft-NMS策略時(shí)間消耗的差距幾乎為零,但是后者檢測精度有小幅度的提升.在各個階段的優(yōu)化之后,最終表現(xiàn)如表3所示.將本文算法與MobileNet-YOLO4算法對比發(fā)現(xiàn),改進(jìn)后的算法時(shí)間消耗差距只增加了0.0007s,速度FPS減少了4.57,但是mAP增長了2.58%,記錄每個實(shí)驗(yàn)的20個目標(biāo)類的檢測平均精度AP如圖8所示.將本文算法的與原始YOLO4算法對比發(fā)現(xiàn),本文模型參數(shù)量只有原模型的1/4,速度FPS提升了26,精度mAP只下降了0.52%.

表3 本文算法優(yōu)化過程中各個階段的模型對比Table 3 Comparison of models in each stage of the algorithm optimization process in this paper

可見本文基于CBAM機(jī)制的MobileNet-YOLOv4實(shí)時(shí)目標(biāo)檢測算法利用MobileNet網(wǎng)絡(luò)中的深度可分離卷積網(wǎng)絡(luò)層的技術(shù),在檢測精度達(dá)到主流水平的同時(shí),檢測速度有了進(jìn)一步的提升,同時(shí)參數(shù)量也大大減小,這有利于部署在計(jì)算能力和內(nèi)存等資源有限的嵌入式設(shè)備上.利用VOC 2007和VOC 2012的訓(xùn)練集進(jìn)行聯(lián)合訓(xùn)練,然后基于PASCAL VOC2007測試集進(jìn)行評估,圖8(a)展示了和該訓(xùn)練網(wǎng)絡(luò)對20個目標(biāo)類的檢測平均精度AP值和總的mAP.

為了進(jìn)一步驗(yàn)證改進(jìn)模型的性能,引入誤檢率MR,并將本文算法與MobileNet-YOLOv4算法進(jìn)行誤檢率比較,如圖8(b)所示,由圖可知,本文算法明顯優(yōu)化了MobileNet-YOLOv4,降低了大部分類別的誤檢率,提高了目標(biāo)檢測的平均檢測精度.

圖8 mAP比較和誤檢率比較Fig.8 mAP comparison and error detection rate comparison

為了更直觀的體現(xiàn)改進(jìn)的目標(biāo)檢測算法的性能,如圖9所示,列舉了本文設(shè)計(jì)的輕量化YOLOv4模型在航拍圖片和視頻上的運(yùn)行結(jié)果.每行分別表示不同的航拍條件,從左到右依次是:正常路面情況、拍攝光照不足、畫面有遮擋、相機(jī)視角傾斜、拍攝停下的車輛、拍攝實(shí)時(shí)視頻.每列分別表示不同的檢測算法,從上到下依次是:拍攝的原始圖像、原始的YOLOv4模型、本文提出的優(yōu)化后的MobileNet-YOLOv4模型.從圖中可以看出,本文提出的模型能夠有效地對圖像、實(shí)時(shí)視頻中所包含不同車輛的大小位置進(jìn)行檢測.相比于原算法來講,檢測結(jié)果并沒有很大差別,但是實(shí)時(shí)性更佳.最后,將本文模型與近3年來目標(biāo)檢測模型在同一編程環(huán)境以及相同數(shù)據(jù)集上進(jìn)行訓(xùn)練,結(jié)果對比如表4所示.

圖9 在不同拍攝條件下,不同算法得到的檢測和分割結(jié)果Fig.9 Detection and segmentation results obtained by different algorithms under different shooting conditions

表4 本文模型與近年來目標(biāo)檢測模型的對比Table 4 Comparison between the model in this paper and the target detection model in recent years

3.3 基于無人機(jī)的飛行試驗(yàn)

圖10 估計(jì)所需的最低幀處理速率Fig.10 Minimum required frame processing rate

實(shí)驗(yàn)使用四旋翼無人機(jī),無人機(jī)飛控為Pixhawk,平臺搭載ZED雙目立體相機(jī),獲取豐富的環(huán)境信息,提高無人機(jī)的智能感知和場景理解能力.在無人機(jī)平臺上,圖像處理模塊要使用體積較小且算力足夠嵌入式平臺,常用的是Nvidia Jetson TX2和Raspberry Pi 4B.

表5 NVIDIA Jetson TX2的系統(tǒng)規(guī)格和軟件Table 5 NVIDIA Jetson TX2 system specs and software

表6 Raspberry Pi 4B的系統(tǒng)規(guī)格和軟件Table 6 Raspberry Pi 4B system specs and software

表7 使用 Nvidia Jetson TX2和Raspberry Pi的FPSTable 7 FPS using Nvidia Jetson TX2

表5和表6展示了這兩種嵌入式系統(tǒng)的規(guī)格.將輕量級模型分別部署到兩個嵌入式系統(tǒng)中,對無人機(jī)航拍圖像進(jìn)行檢測處理,所得到的檢測結(jié)果如圖11所示,檢測速度如表7所示,飛行試驗(yàn)顯示TX2上的FPS達(dá)到了21.8,相比于YOLOv4提高了3.74倍.并且在Raspberry Pi上,FPS也達(dá)到了8.5,故將本文算法部署到無人機(jī)裝載的嵌入式平臺上,能夠?qū)脚囊曇爸械能囕v目標(biāo)進(jìn)行實(shí)時(shí)識別和定位.

圖11 對無人機(jī)采集的視頻進(jìn)行實(shí)時(shí)目標(biāo)檢測的結(jié)果Fig.11 Target detection and insance segmentation on images collected by drones

4 總結(jié)與展望

本文提出的基于CBAM機(jī)制的MobileNet-YOLOv4實(shí)時(shí)目標(biāo)檢測方法,首先將MobileNet替換為YOLOv4的主干網(wǎng)絡(luò),并且利用MobileNet中的深度可分離卷積技術(shù),將YOLOv4中的部分標(biāo)準(zhǔn)卷積替換為深度可分離卷積.接下來優(yōu)化MobileNet- YOLOv4模型,通過嵌入卷積注意力機(jī)制CBAM提高了卷積神經(jīng)網(wǎng)絡(luò)輸出特征圖的全局特征;其次通過引入 Soft-NMS 有效地降低了因?yàn)閭鹘y(tǒng)非極大抑制NMS算法導(dǎo)致的密集物體的相鄰框漏檢問題.最終在PASCAL VOC數(shù)據(jù)集上的測試結(jié)果,表明算法在保證檢測精度的前提下,有效地降低了參數(shù)量和復(fù)雜度,檢測速度有了大幅度的提升.這有利于將算法部署到計(jì)算能力和內(nèi)存等資源有限的無人嵌入式平臺上,使得無人機(jī)能夠?qū)σ曇爸械哪繕?biāo)進(jìn)行實(shí)時(shí)識別,提高無人機(jī)的場景理解能力,廣泛應(yīng)用于無人機(jī)檢測和跟蹤車輛、捕獲違規(guī)行為的智能交通領(lǐng)域以及災(zāi)害救援、土地變化監(jiān)控等應(yīng)用.