基于YOLOv5的輕量化交通標志檢測方法

何鑫，陳輝

(安徽理工大學 計算機科學與工程學院，安徽 淮南 232001)

隨著人工智能技術的快速發展,智能汽車的自動駕駛已經逐步走入人們的生活,并已開始投入商業應用。交通標志識別(traffic sign recognition,TSR)是自動駕駛的基礎性技術,也是自動駕駛系統的關鍵技術[1],對其研究也逐漸成為目標識別領域的熱點。交通標志識別方法可以分為兩大類：傳統識別方法和基于深度學習的識別方法[2]。傳統識別方式主要基于顏色、形狀等方法識別,如Romdhane等[3]結合顏色分割、HOG(histogram of oriented gridients)和SVM分類器以識別交通標志;Shang[4]基于HOG特征訓練級聯分類器的方法識別交通標志。近年來隨著卷積神經網絡的高速發展,深度學習的檢測方法逐漸成為熱門的研究方向。深度學習方法有兩類,一類是先進行候選框的篩選再提取特征進行分類、預測回歸的兩階段目標檢測算法(two-stage),代表性的算法有RCNN[5]、Fast-RCNN[6]、Mask-RCNN[7];另一類是一次檢測即可產生預測的單階段目標檢測算法(one-stage),典型的算法有YOLO系列[8]、SSD系列[9]。雖然傳統依賴顏色、形狀的交通標志識別方法魯棒性較強,但實時檢測中受限于天氣情況、光照條件、交通標志損壞等因素的干擾,基于檢測RoI(region of interest)進行分類的傳統檢測方法往往不能在其他因素影響下達到理想的識別效果,且算法計算量大,而基于深度學習的識別方法特征表達能力強,也不易受到復雜環境的影響而導致識別精度大幅下降。因此,深度學習用于交通標志識別在某些方面更具優勢。如江金洪等[10]將深度可分離卷積((depthwise separable convolution,DSC)[11]引入YOLOv3[12],并更換損失函數MSE為GIOU來檢測交通標志;呂禾豐等[13]替換YOLOv5損失函數為EIOU[14],并使用加權Cluster非極大抑制(NMS)改進YOLOv5中原NMS以提升檢測精度。雖然這些改進算法相較于傳統檢測方法在精度方面有所進步,但依舊存在算法計算量大、模型尺寸大等問題。

綜上所述,基于顏色、形狀和基于深度學習的交通標志識別技術分別存在著檢測精度低和計算量大、模型尺寸大、推理速度慢的問題,使得模型難以部署到移動設備中。鑒于此,本文以YOLOv5為基礎網絡結構,通過GhostNet和坐標注意力機制的有效位置融合以及SIOU損失函數優化使用,提出一種輕量化改進的YOLOv5交通標志檢測模型。

1 基于 YOLOv5的輕量化改進算法

1.1 主干網絡的改進

YOLO算法是Redmon等于2016年首次提出的目標檢測算法,在不同領域得到廣泛應用。YOLOv5是YOLO系列的第五個版本,主要由輸入端、Backbone、Neck、Prediction四部分組成。輸入端主要負責圖像的輸入,Backbone負責對輸入圖片進行不同程度的特征提取,Neck部分融合、組合這些不同尺寸特征提取后的特征圖,最后由Prediction預測。YOLOv5依照深度寬度的不同劃分為五個不同的版本,其中YOLOv5m、YOLOv5l以及YOLOv5x網絡結構復雜、體積過大,難以部署到嵌入式設備中,YOLOv5n模型尺寸過小,精度難以達到一定的高度,故本文選用深度較小,特征圖寬度較小的YOLOv5s作為基礎算法加以改進。改進模型的整體框架如圖1所示。

圖1 改進模型框架

改進模型利用GhostNet網絡對YOLOv5s整體架構進行輕量化,將結構復雜、參數較高的C3模塊與Ghost Bottleneck融合為C3Ghost模塊,替換模型中所有原Bottleneck模塊,剩余Conv模塊更換為Ghost模塊,實現網絡結構的精簡,最后在模型的小目標尺寸特征提取區部署CA坐標注意力機制,即YOLOv5s中Backbone最后一層,增強模型小目標的識別能力。

1.2 GhostNet網絡

鑒于YOLOv5s存在網絡復雜度高的問題,使用構建好的C3Ghost模塊和普通Ghost模塊融合搭建新YOLOv5s網絡框架,不僅能有效減小模型整體參數以壓縮模型體積,且能保證其具有一定的準確性。GhostNet是由Han等[15]在2020年提出的一種輕量化神經網絡架構,該網絡不僅能大幅度減少模型參數,且具有較高的識別性能,可以在同等計算成本下擁有趕超Mobilenetv3[16]的準確性。主流卷積神經網絡在訓練過程中會產生大量特征映射冗余,GhostNet網絡提出一種新的Ghost模塊。圖2為GhostNet網絡結構圖,可以看出,GhostNet主要分兩步,首先進行常規卷積運算得到一半特征圖,再通過廉價的線性變換將其中一半卷積生成具有更多相似特征的另一半特征圖,最后兩部分相結合輸出,這種方式提高了卷積量,同時減少了參數量和計算量。

圖2 GhostNet網絡結構

GhostNet 參考 了ResNet先升維再降維的網絡結構,如圖 3 所示。Ghost Bottlenecks 由兩個Ghost 模塊構成,第一個 Ghost 模塊主要作用是擴展并增加通道數量;第二個 Ghost 模塊主要負責減少匹配通道數量以及鏈接兩模塊輸入輸出。GhostNet 借鑒 MobileNetv2[17]中的設計技巧,在第一個 Ghost 模塊后不使用 ReLu 激活函數,其余層都進行 BN (batch normalization) 批歸一化和 ReLu非線性激活。這種特殊的結構設計,不但能夠壓縮參數量和計算量,還可以優化特征圖,提高模型的性能。

圖3 GhostNet結構圖

圖3所示C3Ghost模塊參考CSP模塊的構建思路,利用兩個CBS(由普通卷積加批歸一化加silu激活函數構成)和一個Ghost Bottlenecks組成。輸入特征信息分別通過Ghost Bottlenecks以及CBS處理拼接成新的特征圖,這樣的構成方式減少了整體參數,同時還能保證特征信息的最大保留。

1.3 CA坐標注意力

對于輸入圖像來說,背景復雜,交通標志空間占比小,模型往往會受到大量冗余背景信息干擾導致損失小目標的特征信息,致使識別精度下降。此外,經過輕量化后的YOLOv5s中存在大量Ghost模塊,Ghost模塊使用分組卷積,減少了通道的數量,使得通道特征僅與自己建立聯系。這一方面令參數量降低、網絡復雜度減小,但是另一方面會影響對全局特征的提取,導致交通標志檢測精度降低。因此本文將位置信息嵌入到通道注意力的坐標注意力機制[18],引入CA注意力機制到已優化的YOLOv5s網絡結構Backbone部分的最后一層的C3模塊,幫助模型把握全局特征,抵消部署Ghost模塊帶來的特征缺失。

有別于其他主流注意力模塊,如SE通過全局池化來計算通道注意力[19],CA將全局池分解為兩個一維的特征編碼以緩解二維池化造成的位置信息缺失,并使坐標信息和輸入特征圖形成交互。圖4為坐標注意力機制CA的網絡結構圖,由坐標信息的嵌入和坐標注意力生成兩部分構成。

圖4 CA網絡結構圖

首先,坐標信息的嵌入,對于給定維度為C×H×W的輸入特征圖,使用空間范圍分別為(H,1)或(1,W)的池化層沿X軸和Y軸兩個不同方向坐標的每個通道進行編碼,經過編碼平均池化后得到尺寸為C×H×1和C×1×W的特征圖;其次,坐標注意力生成,通過拼接和1×1卷積將兩種尺寸的特征圖聚合成一個C/r×1×(W+H)的特征圖,維度也從C維降成了C/r維;再分別用1×1卷積進行維度提升,并使用Sigmoid激活函數激活;最后進行特征圖重加權,即將生成的兩個同緯度特征圖與輸入特征圖相乘,完成對模型坐標注意力的添加。

圖5(a)為主干網絡沒有添加CA的識別效果,5(b)為添加后的效果,可以看出未添加CA會出現對小目標漏檢的情況,添加后漏檢情況得到改善。

(a)未添加CA (b)添加CA圖5 CA識別效果對比圖

1.4 邊框回歸損失函數

損失函數能夠度量模型的預測值與真實值差異程度,其設定也會直接影響模型優化的速度和定位的準確性,因此選擇合適的損失函數對模型起著至關重要的作用。

YOLOv5原損失函數為CIOU損失函數[20],CIOU不僅吸收了GIOU[21]、DIOU[20]的優點,即在IOU的基礎上加上對預測框和真實框的重疊面積、中心點距離的考慮,而且提出對回歸框寬高比進行損失計算,使得目標框回歸更加趨于穩定。

雖然CIOU考慮到了預測框和真實框的重疊面積、中心點距離以及寬高比,但忽略了其方向之間不匹配的問題,導致模型收斂速度不理想。本文選用SIOU損失函數[22]替換YOLOv5中的CIOU損失函數。SIOU吸收了CIOU的優點,并添加了對真實框和預測框之間角度的損失計算,提高了模型的訓練效率和收斂速度。SIOU損失函數由角度損失(angle cost)、距離損失(distance cost)、形狀損失(shape cost)、IOU損失(IOU cost)四部分組成。

角度損失的添加是為了最大程度地減少與距離相關的變量數量,在模型訓練收斂過程中首先將預測帶到最靠近的坐標軸(X軸或Y軸),再沿此軸繼續進行訓練,其計算公式為

(1)

式中：ch為真實框和預測框的高度差,σ為真實框和預測框中心點的距離。

考慮到角度損失的添加,需要實現角度損失和距離損失的平衡。距離損失公式為：

(2)

(3)

形狀損失指的是兩框中心位置的偏差,計算公式為：

(4)

(5)

式中：w和wgt分別為預測框和真實框的寬度,h和hgt分別為預測框和真實框的高度,θ體現了對形狀損失的關注程度。

IOU損失表示真實框和預測框之間的交集和并集之比,計算公式為

(6)

最終損失函數表達式為

(7)

2 TensorRT

TensorRT是由NVIDIA公司提出的一個基于深度學習的高性能學習推理優化器,它支持TensorFlow、Caffe、Mxnet、Pytorch等幾乎所有的深度學習框架,將TensorRT和NVIDIA的GPU結合起來,能在幾乎所有的深度學習框架中進行快速和高效的部署推理。本文利用TensorRT對已訓練好的模型權重進行優化,以加速模型推理速度。

3 實驗結果與分析

3.1 數據集

為了驗證輕量化改進模型的效果,本文采用由長沙理工大學采集并制作的交通標志數據集CCTSDB2021(csust Chinese traffic sign detection benchmark 2021)[23]作為實驗所用數據集,驗證本文算法。CCTSDB2021由三大類交通標志構成,分別為指示標志、禁止標志和警告標志,數據集制作規范符合實驗要求。隨機選取其中6 000張圖片,按照VOC2007數據集格式調整,按照4∶1的比例劃分為訓練集和驗證集進行實驗驗證。

3.2 數據增強處理

為了提升模型的泛化性能,使用Mosaic和Mixup數據增強方法處理數據集。Mosaic數據增強方法通過隨機排列拼接經過裁剪、縮放的四張輸入圖片形成新的圖片輸入訓練模型,擴充小樣本目標的同時提升模型訓練速度,增強網絡的魯棒性;而Mixup數據增強方法則是通過隨機疊加融合兩張不同的圖片豐富樣本數量,實現對小目標檢測能力和模型魯棒性的增強。圖6為使用Mosaic和Mixup數據增強后的效果圖。

(a)Mosaic (b)Mixup圖6 數據增強后的效果圖

Mixup計算公式為：

(8)

(9)

式中：(xi,yi)和(xj,yj)是原始數據的隨機樣本對;λ是服從beta分布的一個參數,λ∈[0,1]。

為驗證數據增強方法對YOLOv5s模型檢測產生的影響,將未開啟數據增強與開啟數據增強的YOLOv5s模型進行實驗對比,表1為對比實驗結果。

表1 數據增強前后對比結果

從表1可以看出,采用數據增強后的YOLOv5s模型MAP值有了提升,說明在YOLOv5s模型中,上述兩種數據處理方式均可以有效提升模型的平均精度。

3.3 實驗配置

本實驗使用Pytorch深度學習框架搭建、訓練改進模型,實驗環境為Ubuntu20.04,CPU采用Intel(R)Xeon(R)Platinum8255C,GUP為V-100SXM2,內存為32 G。實驗輸入的圖像尺寸為640×640,批訓練處理量為32,初始學習率為0.01,IOU閾值為0.2,權重衰減為0.000 5,共訓練150輪。

3.4 評價指標

實驗使用平均精度(mean average precision,MAP)、推理速度、參數數量、模型大小作為對改進模型的評價指標。其中MAP通過準確率(Precision)和平均正確率(average precision,AP)計算得出。

準確率即網絡檢測到正樣本數量占總樣本數量的比率,公式為

(10)

式中：TP(true positive)為被正確識別為正樣本的正樣本數量,FP(false positive)為被錯誤識別為正樣本的負樣本數量。AP、MAP分別為預測樣本的平均正確率、平均精度,所有類別的平均精度為：

(11)

(12)

3.5 實驗結果與分析

本文選用YOLOv5s、YOLOv3-tiny、YOLOv4-tiny、YOLOv5s+Shufflenetv2、SSD五種網絡進行交通標志檢測實驗,選用模型涵蓋了近年來目標檢測領域內經典框架和輕量型改進算法,并將本文模型與之進行性能對比,以驗證本模型的可行性,對比結果見表2。

表2 不同算法對比結果

從表2可以看出,本文模型與輕量化網絡模型YOLOv3-tiny相比平均精度提升了6.7%,參數量和模型大小分別減少了63.8%和61.9%;與另一主流輕量化網絡模型YOLOv4-tiny相比平均精度提升了3.7%,參數量降低了46.9%,模型大小降低了71.8%;和引入Shufflenetv2的YOLOv5s的輕量化模型YOLOv5s+Shufflenetv2相比精度有較大提升,參數量和模型大小均有所減小;相較于主流檢測模型SSD,精度大幅提升,參數量和模型大小大幅減小;與原模型YOLOv5s相比,平均精度只有0.1%的損失,但參數量和模型大小大幅度減小。

改進YOLOv5和原模型YOLOv5s對交通標志數據集訓練結果的平均精度如圖7所示,可以看出改進模型與原模型相比平均精度幾乎沒有差距,在訓練前期效果甚至略強于原模型。

圖7 平均精度對比圖

另外改進模型推理速度也有提升。實驗結果表明,原模型的推理時間為9.4 ms,改進模型的為9.2 ms,提升幅度為2%;在經過TensorRT加速后,模型推理時間更是可以達到4 ms,比原模型提升了57.4%。

改進模型對比主流檢測模型、輕量化網絡模型,不僅在精度上更具優勢,且參數量和模型大小也遠遠小于其他模型;在和原模型的對比上,雖然精度有些許損失,但模型尺寸和參數量減小,且推理速度遠超原模型,綜合來看改進模型的綜合性能優于原模型。

為了驗證所改部分的有效性,對改進模型進行消融實驗。實驗對比結果見表3。表3中改進模型1為融入GhostNet的YOLOv5s模型,改進模型2添加了CA注意力。從表3可以看出,利用GhostNet對YOLOv5網絡架構進行優化使得參數量和模型尺寸減小,但模型精度過低;而在主干網絡加入CA坐標注意力模塊后,模型對重點信息定位能力有所加強,精度增加較大,且參數量和模型尺寸進一步減小;YOLOv5s將原損失函數CIOU優化為SIOU后,彌補了CIOU在真實框和預測框角度問題的不足,精度略微增加;可以看出本文模型在保證了較少精度損失的同時,有效地壓縮了參數量和模型大小。

表3 消融實驗對比結果

4 結論

1)融合CA注意力與GhostNet網絡結構的YOLOv5模型,相較于原模型,參數量減少50%,模型大小減小53.5%,對嵌入式設備的有限計算資源占用較少,利于部署。

2)輕量化模型精度損失僅0.1%,并未因參數量、計算量的減少造成過大的精度損失,相對而言改進模型綜合性能更強。

3)改進模型推理速度快于原模型,相較于原模型推理速度提升了57.4%。