基于改進的YOLOv3-SPP算法目標檢測研究

洪畢輝,李文彬,朱 煒,王曉鳴,張克斌

(南京理工大學 智能彈藥技術國防重點學科實驗室, 南京 210003)

0 引言

當今,隨著深度學習技術的快速發展,在現代戰場上智能彈藥扮演著越來越重要的角色。從廣義上講,智能彈藥是利用嵌入式系統將相關技術程序應用于無人機等移動設備上,裝載相關彈藥武器,能夠在發射后實現簡易具備察/打一體化能力,并通過一定的智能化技術完成命中目標的彈藥武器。在執行相關任務中,關鍵的一點是利用機載攝影系統拍攝前方實時戰況,對目標進行識別檢測和定位,并及時將情況反饋給指揮人員。然而,戰場上利用無人機對目標進行檢測仍然面臨著許多困難挑戰:機載智能平臺設備算力低下;無人機屬于低、小、慢飛行器,容易受到周圍環境干擾,其所采集到的目標容易受到環境光線和拍攝角度等影響。

目前基于深度學習的目標智能識別檢測技術已成為新型智能彈藥目標檢測的主流技術中重要的一環。深度學習(Deep Learning)的概念由Hinton教授于2006年在《Science》提出,自此人類對于神經網絡的研究進入了深度學習時代[1]。基于深度學習的目標檢測框架主要分為兩類,一類是基于Two-Stage方法,代表主流框架有R-FCN[2]、MaskR-CNN[3]、Fast[4]/Faster R-CNN[5]等,將檢測任務分為回歸和分類;另一類是基于One-Stage方法,代表的主流框架有YOLOv3-SPP[6]、YOLO9000[7]、YOLOv3[8]、SSD[9]等,同時完成檢測和回歸任務。2種方法各有所長,Two-Stage任務準確率比較高,但是相對的速度較慢,One-Stage能夠達到實時性,然而在一定程度上犧牲了精度,董文軒等[10]得出One-Stage雖然在推理速度大幅提高,但是檢測精度依舊是改進的主要內容。

為了提高智能彈藥目標檢測技術的算法性能,本文建立了一個包含多種地面軍事作戰單元的小型數據集,對數據集使用數據增強處理手段,使數據集更加貼近于軍事作戰環境,進一步增加了數據樣本數量以提高模型的魯棒性;基于One-Stage方法下的YOLOv3-SPP目標檢測網絡基礎框架,采用K-means++[11]聚類算法提取數據集中適用于本文數據集的錨框,改善檢測結果的最終誤差;將DIoU[12]和Focalloss[13]分別作為預測框的定位損失函數和分類損失函數來降低正負樣本對模型的影響,進一步提高算法的檢測結果。實驗結果表明本文提出基于改進的YOLOv3-SPP算法對復雜戰地環境下軍事目標具有更好的檢測精度。

1 小型目標數據集建立

地面武器系統,是智能彈藥打擊的主要目標,為此本文選取地面目標作為首要打擊目標,建立數據集,對YOLOv3-SPP算法進行改進研究。

根據陸地戰場環境分析,整體軍事行動單元可大致分為:以坦克為主的容納攻防一體化的突擊武器;以裝甲車為主的供步兵機動作戰用的裝甲戰斗車輛,用于協同坦克作戰,快速機動步兵分隊;以雷達為主的無線電波探測設備,能全天候工作是具備直接威脅的后方軍事目標。然而,在真實作戰環境中,通常會通過對軍事作戰單元采用各種偽裝手段來提升生存能力,減低被探測到幾率。

針對復雜戰地環境下的整體軍事行動單元,從互聯網上下載多種復雜環境下目標圖像,建立一個小型軍事目標數據集。本文中所采用的小型軍事目標數據集包括坦克、步戰車、雷達以及為最后實物驗證用普通小轎車等4類目標圖像。數據集部分目標圖像如圖1所示。將數據集中各種軍事目標按照8∶1∶1的比例劃分為訓練集、測試集、驗證集。為方便后期進行深度學習訓練,再利用開源軟件LabelImg對數據集進行圖像圖形注釋,將所需檢測目標用方框注釋。針對數據集將檢測目標分為坦克、坦克炮管、步兵戰車、雷達、轎車五類,其中坦克炮管作為坦克打擊首要目標,進行專門的標注并識別。數據集注釋以PascalVOC格式保存為XML文件,該文件包含了圖像尺寸、目標類別及坐標等信息。

圖1 數據集部分示例圖像Fig.1 Sample images for the dataset

本文采用有監督學習單樣本數據增強處理方法,增加訓練樣本的數量以及多樣性(噪聲數據),防止過擬合,提升模型魯棒性。為了使數據集更加貼近于實際戰地環境,兼顧無人機在前方可能會遇到各種不確定因素,針對軍事目標會以任何角度任何方向出現在圖片中,對樣本進行鏡像處理以增加樣本存在更多角度;為提高在惡劣電磁對抗環境下的生存能力[14]采用高斯噪聲對單樣本處理。除此之外較為新穎的增強方式還有隨機擦除(Random Erasing)[15]和Cutout[16]。

通過對軍事目標數據集進行數據增強處理,將數據量擴增為原數據集的4倍。數據集具體信息明細如表1所示。

表1 數據集信息明細Table 1 Dataset information breakdown

2 YOLOv3-SPP目標檢測算法改進

YOLO從2016年被Redmon等[17]提出開始,至今已發展到YOLOv5,其中以YOLOv3最為經典。YOLOv3主要分為特征提取和特征預測兩部分。特征提取主體是一個Darknet-53網絡,其大量使用殘差(Residual Network)跳躍連接,去掉全連接層。為了降低池化帶來的梯度負面效果,使用卷積下采樣代替池化層,使YOLOv3成為一個全卷積網絡,可以支持不同大小輸入,支持全圖end-to-end訓練,更好的學習context信息。YOLOv3使用3種不同尺度對特征進行預測,并在3種尺度(13×13、26×26、52×52)上輸出對物體的預測。與YOLOv3相比,YOLOv3-SPP是在第五、六層卷積之間加入一個SPP模塊。該模塊主要由3個最大池化層(5×5、9×9、13×13)和1個跳層連接組成,具體如圖2所示。SPP模塊借鑒了空間金字塔的思想,提取局部特征和全局特征,可以提升模型的感受野,有利于待檢測數據集中目標大小差異較大的情況。在YOLOv3-SPP原始模型中,采用了MSELoss和BCEWithLogitsLoss兩種損失函數,前者作為損失函數不能很好的區分預測框和真實框的位置關系,后者作為損失函數難以解決正負樣本不平衡問題。因此,本文在YOLOv3-SPP框架基礎上,改進錨框設計方法,將DIoU和Focalloss替換均方誤差函數和交叉熵函數,進一步改善軍事目標檢測算法精度。

圖2 SPP模塊結構圖Fig.2 SPP module structure diagram

2.1 錨框設計

YOLO系列自YOLOv2開始引入錨框(Anchorbox)機制,使錨框的生成不和以前一樣靠經驗所得,而是根據數據集中目標的真實框(Ground Truthbox)位置,使用K-means維度聚類[18]得到。因此算法僅需要在錨框的基礎上進行微調即可得到與目標物體有良好的匹配的預測框(Predictionbox),這樣模型可以更快的收斂,檢測效果會更好。K-means聚類算法的思想就是對于給定的樣本集,按照樣本之間的距離大小,將樣本集劃分為K個簇,隨機選擇K個聚類中心。簇內的點盡量緊密的連在一起,簇間的距離盡量的大。YOLOv3-SPP默認使用的錨框是在COCO數據集上使用K-means聚類得到了9個anchors。COCO數據集全稱Microsoft Common Objects in Context,該數據集一個大型的、豐富的物體檢測,分割和字幕數據集。圖像包括91類目標,328 000影像和2 500 000個label。目前為止有語義分割的最大數據集,提供的類別有80類,有超過33萬張圖片,其中20萬張有標注,整個數據集中個體的數目超過150萬個[19]。可以說,默認使用的9個Anchors具有一定的適用性,但直接使用該錨框尺寸并不完全適用于本文的小型數據集,并且K個聚類中心位置的選擇對最后的聚類結果和運行時間都有很大的影響。因此,為了更好的選取適用于本文數據集的錨框尺寸,選取K-means++算法來代替原有K-means算法。K-means++算法在K-means基礎上優化了對聚類中心的選取,摒棄原來隨機選擇K個聚類中心的思想,改用隨機從數據點集合里選擇一個聚類中心,通過使用輪盤法[20]取到其余聚類中心。K-means++算法計算步驟如下:

步驟1從數據點中隨機選取一個數據點作為初始聚類中心u1。

步驟4重復步驟2和步驟3,直到選出k個中心(j=k)。

步驟5在現有聚類中心基礎上,繼續使用K-means聚類。

最終得到9個錨框的尺寸分別為:(28,6),(80,11),(67,30),(76,73),(149,43),(129,111),(224,79),(301,129),(308,182)。

2.2 改進損失函數

YOLOv3-SPP損失函數Lall由定位損失函數Lbbox、置信度損失函數Lconf與分類損失函數Lcls組成,如式(1)所示:

Lall=Lbbox+Lconf+Lcls

⑴

MSELoss,也叫均方損失函數,用在做定位損失任務中。該函數的特點是光滑連續可導,方便用于梯度下降,但是MSE計算公式是模型預測值f(x)與樣本真實值y之間距離平方的平均值,因此,離得越遠,誤差也就越大,受離群點影響很大。MSE計算如式(2)所示:

(2)

MSE用于做定位損失函數,需要計算預測框寬w、高h、中心點x、y的損失。這四者是需要解耦獨立的,但實際上這四者并不獨立,因此,將IoU系列損失函數用來能捕捉到它們之間的關系,IoU也叫交并比,是深度學習目標檢測中常用的衡量指標,計算如式(3)所示:

(3)

式中:AP為預測框面積,AG為真實框面積。

它可以用來評價預測框和真實框的距離,很好地反映預測檢測框和真實檢測框的檢測效果,擁有對尺度不敏感特性。但作為損失函數用來計算定位損失,會存在2個框沒有相交的情況,根據式(3),IoU為0,不能反映兩者的距離,并且因為損失為0,導致梯度沒回傳,神經網絡訓練收斂速度很慢,無法精確的反映兩者重合度情況。為了避免這種情況,本文引入DIoU(DistanceIoU)損失函數作為定位損失函數。計算公式如(4)所示:

(4)

式中:ρ表示預測框和真實框之間的歐氏距離,aP和aG表示2個預測框和真實框的中心,c表示預測框和真實框閉包區域的對角線距離。DIoU損失函數設計了2個框之間的距離ρ作為懲罰項,使得檢測效果更符合實際。c的作用是防止損失函數的值過大,加快神經網絡訓練過程,提高模型收斂速度,得到更準確的預測框回歸結果。

BCEWithLogitsLoss和BCELoss,一組經常用到的二元交叉熵損失函數,常用于二分類問題,前后者區別在于后者的輸入為經過sigmoid處理過的值。交叉熵損失函數計算如式(5)所示。在模型中,采用多尺度特征輸出對目標做檢測任務,模型輸出包括13×13、26×26、52×52三種檢測特征圖,特征圖中一個像素點為一個單元格,每個單元格會有3種檢測框,因此會輸出 10 647個檢測框。但只有極少部分檢測框會包含物體,所以大部分的候選框都會被認為是負樣本。這種情況在檢測小物體時會更加明顯。因此在訓練過程中正負樣本數量差異過大,容易導致過多的負樣本在訓練過程占據主導作用,使網絡模型無法得到有用信息,降低了檢測準確率。

(5)

為了解決正負樣本不平衡問題,引入Focal Loss函數。該函數基于交叉熵函數改進而來,引入加權因子α抑制正負樣本數量的失衡,引入γ控制簡單/難區分樣本數量失衡。計算公式如(6)所示:

(6)

綜上得,YOLOv3-SPP損失函數Lall如下:

Lall=Lbbox+Lconf+Lcls=

3 實驗結果分析

3.1 實驗環境配置與檢測結果

將本文改進的YOLOv3-SPP算法在小型軍事目標數據集上進行訓練和測試,以驗證其有效性。本文實驗所使用的環境配置:實驗均在Windows操作系統上進行,中央處理器為Inteli7-10750H,圖形處理器為NVIDIA GeForce RTX2060,顯卡內存為6GB,采用Pytorch+TensorFlow為基礎框架,使用Python語言進行編程。

本實驗采用平均精度(mean Average Precision,mAP)、準確率(Precision,P)、召回率(Recall,R)來評價模型檢測的綜合性能。平均精度是用來衡量訓練模型在檢測類別上的好壞程度,準確率表示預測為正的樣本中有多少是真正的正樣本,召回率表示樣本中的正例有多少被預測正確了。定義分別表示為

其中:N為測試樣本個數,p(n)為檢測n個測試樣本的精度值,TP表示為檢測正確目標個數,FN表示為未檢測出來的目標個數,FP表示為檢測錯的目標個數。

模型訓練過程:在模型訓練過程中,使用隨機梯度下降法(SGD)進行模型訓練,受到顯存的影響,Batch-size大小設置為2,每迭代32步更新一次參數,初始學習率設置為0.001,權重衰減系數和動量因子分別設置為0.000 5和0.93,訓練批次設置為200次。應用改進軍事目標檢測算法(YOLOv3-SPP-DFDA+)對復雜戰地環境下軍事目標小型數據集進行訓練檢測,其檢測效果如圖3所示。

圖3 標檢測算法檢測結果示例Fig.3 Example of an object detection algorithm result

3.2 檢測結果分析

圖4展示的是模型在訓練過程中的曲線變化,其中圖4(a)表示檢測精度,圖4(b)表示平均精度,圖4(c)表示召回率,圖4(d)、圖4(e)分別為訓練損失和驗證損失。黑色線代表原模型(YOLOv3-SPP,模型A)用無增強處理數據集訓練的損失曲線;紅色線代表模型(YOLOv3-SPP-DFDA,模型B)替換成DIoU函數和Focalloss函數(DF)后并對數據集進行數據增強處理(DA)的損失曲線;藍色線代表模型(YOLOv3-SPP-DFDA+,模型C)引進K-means++錨框后的損失曲線;綠色線代表對比模型(YOLOv4-SPP,模型D)用增強處理數據集訓練的損失曲線。橫坐標表示訓練批次,縱坐標表示損失值。

圖4 模型訓練過程曲線對比Fig.4 Comparison of model training process curves

通過對圖中曲線的對比可以得出:

模型A→模型B:模型總體性能有了很大的優化,數據集的增強處理能夠加快訓練模型的收斂速度,增強模型的魯棒性和泛化性,替換成DIoU和Focal Loss函數后進一步提高了模型的定位能力和分類能力,使模型檢測精度提高了近3%,模型的召回率和mAP分別提高了近8%和10%,說明了模型對正確分類和能找到檢測出來目標的能力都有所提高。

模型B→模型C:針對本文所創建的軍事目標數據集使用K-means++得出的錨框系列,能夠使模型預測框的回歸效果更好,提高了模型的回歸精度,在召回率和平均精度近乎不變的情況下,使模型的精確度提高了近6%,使得模型總體對目標的檢測能力有所上升,可以更快的使模型收斂,檢測效果更佳。

模型C→模型D:將文中改進后模型C與經典YOLOv4-SPP模型對比,可以得出模型C在訓練精度、平均精度和準確率上均比YOLOv4-SPP訓練結果來的更好,改進后的模型C更加適用于本文中建立的數據集。

實驗證明,通過上述改進方法,可以有效提高模型的檢測能力。

4 結論

本文中基于YOLOv3-SPP算法,將數據集進行數據增強處理,訓練結果表明數據增強能有效加快模型收斂速度,增強了模型的魯棒性和泛化性;針對預測框和真實框的位置關系、正負樣本不平衡問題,將DIoU和Focal Loss替換原模型中的均方誤差和交叉熵函數,訓練結果表明改進函數能進一步提高了模型的定位能力和分類能力;使用K-means++聚類錨框算法替換原K-means聚類錨框算法,優化了對聚類中心的選取,得到更適用于本文數據集的錨框尺寸,訓練結果表明對預測框的回歸效果更好,提高了回歸精度。

此外,本文數據樣本及類別數量和種類都存在不足問題,并且模型訓練得到的權重比較大,在后續的研究中,擴充數據集和對模型進行輕量化處理都將是本文研究的重點。