基于STN-GhostNet 的輕量化交通標(biāo)志識別模型

張成涵宇，林鈺哲

(杭州電子科技大學(xué) 電子信息學(xué)院，浙江杭州，310018)

0 引言

隨著自動駕駛相關(guān)技術(shù)的不斷發(fā)展，交通標(biāo)志的快速、準(zhǔn)確識別逐漸成為了深入發(fā)展自動駕駛系統(tǒng)的重要部分。其通過對車輛行駛過程中拍攝的交通標(biāo)志圖像進(jìn)行識別和分類，可以為自動駕駛系統(tǒng)的進(jìn)一步操作判斷提供一定的指示或警告。目前隨著深度學(xué)習(xí)技術(shù)的不斷發(fā)展，已經(jīng)出現(xiàn)了一系列運用YOLO 模型實現(xiàn)交通標(biāo)志識別的方法[1,2]。但由于YOLO 系列本身模型尺寸較大，模型本身參數(shù)量較多，導(dǎo)致其識別交通標(biāo)志的速度較慢，這將對自動駕駛的安全性產(chǎn)生一系列的負(fù)面隱患。為了保證實際交通標(biāo)志識別的時效性與精確性，本文提出了一種新型的輕量化模型。通過將Spatial Transformer Networks(STN)網(wǎng)絡(luò)模塊以及GhostNet 網(wǎng)絡(luò)模塊結(jié)合來替換原先YOLOv5s 的主干網(wǎng)絡(luò)，并對其進(jìn)行剪枝操作來實現(xiàn)輕量化處理，可以在保證原YOLOv5s 網(wǎng)絡(luò)識別精度的前提下大幅提高交通標(biāo)志圖像的識別速度，減小模型尺寸，從而便于后續(xù)部署在自動駕駛系統(tǒng)中。

1 交通標(biāo)志識別模型構(gòu)建

■1.1 GhostNet 模型

目前，一些科研工作者認(rèn)為神經(jīng)網(wǎng)絡(luò)中豐富甚至具有冗余信息的特征圖對于其能夠充分理解輸入的圖像數(shù)據(jù)至關(guān)重要。因此，為了使模型能夠以輕量高效的方式實現(xiàn)圖像的特征提取，華為諾亞方舟工作室提出了一種新型的端側(cè)神經(jīng)網(wǎng)絡(luò)架構(gòu)GhostNet[3]。其通過引入一種新型的Ghost 模塊來替代原先的卷積層，以一種代價低廉的線性操作來獲取特征圖信息，從而可以用更少的參數(shù)生成更多的特征。模塊具體的運行流程如下圖所示：

圖1 Ghost 模塊運行流程

Ghost 模塊主要對原先的卷積操作進(jìn)行了改進(jìn)。首先通過主卷積來生成基本的特征圖集合：

其中，X 為輸入特征圖，f'為卷積層，Y'為卷積層卷積后得到的特征圖集合。

最后再將得到的Ghost 特征圖與原先的特征圖進(jìn)行拼接作為最終輸出的特征圖集合Y。相較于直接進(jìn)行卷積得到的特征，用Ghost 模塊進(jìn)行特征提取可以降低模塊所需參數(shù)總數(shù)以及計算復(fù)雜度，同時還能夠增加得到的特征信息量。

之后在其主干網(wǎng)絡(luò)中進(jìn)行殘差塊的設(shè)計，使其可以有效避免出現(xiàn)梯度消失與梯度爆炸的問題，從而提升網(wǎng)絡(luò)的性能。具體的結(jié)構(gòu)如圖2 所示。

圖2 GhostNet 瓶頸層

因此將原先YOLOv5s 的主干網(wǎng)絡(luò)替換為GhostNet，可以有助于提高網(wǎng)絡(luò)模型對輸入交通標(biāo)志圖像的特征提取能力，并且更為輕量高效。

■1.2 STN-GhostNet 算法模型

由于在車輛行駛的過程中，攝像頭拍攝到的交通標(biāo)志難免會出現(xiàn)顯示傾斜或者圖片噪聲干擾等問題。為此需要在GhostNet 的基礎(chǔ)上加入STN 模塊[4]，賦予網(wǎng)絡(luò)基于特征圖本身進(jìn)行空間變換的能力，保證模型實際的識別精度不會受到交通標(biāo)志角度傾斜等因素的干擾。STN 模塊的結(jié)構(gòu)如圖3 所示。

圖3 STN 模塊基本結(jié)構(gòu)

STN 模塊內(nèi)部主要分為Localisation Net、Grid Generator 和Sampler 三部分。通過Localisation Net 對輸入的特征圖進(jìn)行參數(shù)的提取與變換后，利用Grid Generator 對原圖片的像素點坐標(biāo)進(jìn)行確定，得到原圖片與輸出圖片的坐標(biāo)仿射關(guān)系，最后通過Sampler 進(jìn)行像素的采集與轉(zhuǎn)換，對輸入的圖像進(jìn)行旋轉(zhuǎn)調(diào)整，從而使模型具有旋轉(zhuǎn)不變性。保證在不同的拍攝角度下模型依舊可以正常識別不同的交通標(biāo)志，增強(qiáng)模型的魯棒性與泛化能力。

由于STN 模塊本身也可以作為一個layer 加入網(wǎng)絡(luò)的任意位置中，且其計算量很少，對模型進(jìn)行交通標(biāo)志的識別速度基本沒有什么影響。故本文將STN 模塊放置于前面來對輸入的圖片特征進(jìn)行轉(zhuǎn)變后再利用GhostNet 對圖片內(nèi)的特征進(jìn)行提取，從而可以減少網(wǎng)絡(luò)訓(xùn)練中的整體損失，提高模型的識別精度。最后得到的STN-GhostNet 模型結(jié)構(gòu)如圖4 所示。

圖4 STN—GhostNet 結(jié)構(gòu)

將STN-GhostNet 網(wǎng)絡(luò)修改為YOLOv5s 的主干網(wǎng)絡(luò)后完成初步的訓(xùn)練，即可進(jìn)行后續(xù)的輕量化操作。

■1.3 模型輕量化

由于神經(jīng)網(wǎng)絡(luò)整體實際存在著大量冗余的神經(jīng)元參數(shù)，對模型的識別速度會產(chǎn)生較大的影響。為了提高交通標(biāo)志的識別速度，減少模型尺寸大小以適應(yīng)自動駕駛系統(tǒng)的輕量化需求。本文決定使用輕量化處理中的剪枝技術(shù)，具體為結(jié)構(gòu)化剪枝中的通道剪枝來對交通標(biāo)志模型中的冗余神經(jīng)元進(jìn)行篩除。

在進(jìn)行交通標(biāo)志識別模型的剪枝操作前還需對其進(jìn)行稀疏訓(xùn)練。稀疏訓(xùn)練是指在識別出對交通標(biāo)志識別結(jié)果影響較小的通道的同時，去除模型ВN 層中縮放因子值γ較小的通道，使最終剪枝后的模型在尺寸上更加緊湊且計算量有所減少。雖然在進(jìn)行剪枝后可能會降低部分性能，但這些損失可以通過對模型的參數(shù)進(jìn)行微調(diào)來補(bǔ)償，從而使最后得到的模型能夠基本達(dá)到與未經(jīng)過剪枝的網(wǎng)絡(luò)相當(dāng)?shù)木萚5]。具體的計算公式如下所示：

其中(x,y)表示訓(xùn)練時模型的輸入以及目標(biāo)結(jié)果，W 表示可進(jìn)行訓(xùn)練的神經(jīng)元權(quán)值，g(.)是指訓(xùn)練時的稀疏度懲罰，λ 則是對公式的平衡處理。本文在實現(xiàn)對YOLOv5s 網(wǎng)絡(luò)的稀疏訓(xùn)練后，模型尺寸有了較大的改善。此時針對整個ВN層的輸入輸出公式則轉(zhuǎn)變?yōu)椋?/p>

其中Zin 和Zout 分別為ВN 層的輸入?yún)⒘亢洼敵鰠⒘?，?指代minibatch 的大小，這樣的改變可以基本保持網(wǎng)絡(luò)的參數(shù)量不會發(fā)生較大變化。

由于模型的寬度決定了得到特征的豐富程度，如果實際僅選擇模型的部分維度進(jìn)行修剪，則會對交通標(biāo)志識別模型的輕量化效果產(chǎn)生負(fù)面影響。故本文共選擇了三個維度—filter、channels 和filter shapes 來進(jìn)行交通標(biāo)志識別模型的剪枝操作。同時本文還利用文獻(xiàn)[6]中提出的剪枝率多項式問題求解得到這些維度剪枝率的分配結(jié)果，即：

2 來稿首頁注腳應(yīng)注明課題來源及課題名稱，獲基金資助產(chǎn)出的文章應(yīng)以“基金項目：”作為標(biāo)識注明基金項目名稱，并在圓括號內(nèi)注明其項目編號，并附第1作者簡介：姓名(出生年-)，性別(民族)，籍貫，職稱，學(xué)位，研究方向，E-mail。

其中f 用于去掉某個filter；c 用于去掉所有filter 的某個channel；fs 用于調(diào)整filter 的形狀，從而最終達(dá)到模型輕量化的效果。

經(jīng)過輕量化后的交通標(biāo)志識別模型在保證原有的識別精度的前提下，模型的識別速度有了大幅度的提高，模型尺寸也有了較大幅度的降低，可以有效滿足自動駕駛系統(tǒng)的輕量化需求，促進(jìn)后續(xù)在自動駕駛中的實際應(yīng)用。

2 交通標(biāo)志圖像識別實驗與分析

■2.1 交通標(biāo)志數(shù)據(jù)集

本文針對交通標(biāo)志識別的研究實驗數(shù)據(jù)集來自于清華大學(xué)與騰訊發(fā)布的公開交通標(biāo)志數(shù)據(jù)集TT100K。其內(nèi)置有大量的道路場景實拍圖片，為減少模型訓(xùn)練的負(fù)擔(dān)，本文對其進(jìn)行初步篩選，共選出43種基本的交通標(biāo)志進(jìn)行識別分類。其中模型訓(xùn)練的訓(xùn)練集圖片共6598 張，測試集圖片共970 張，驗證集圖片共1889 張。部分?jǐn)?shù)據(jù)集圖片如圖5 所示。

圖5 TT100K 部分交通標(biāo)志圖

■2.2 模型訓(xùn)練

2.2.1 訓(xùn)練環(huán)境與模型超參數(shù)設(shè)置

實際的交通標(biāo)志識別代碼運行環(huán)境為Anaconda 所創(chuàng)建的虛擬環(huán)境，其內(nèi)部主要應(yīng)用了pytorch1.7.1、opencv-python、numpy 以 及matplotlib 等代碼庫來搭建實際的交通標(biāo)志識別模型，并使用cuda11.7 來加速模型的訓(xùn)練過程。PC 端運行硬件則使用了Inter Core i7-10875H(2.3GHz)CPU 以及NVIDIA GeForce RTX 2060 來進(jìn)行模型訓(xùn)練。

同時模型訓(xùn)練時加載了YOLOv5s 預(yù)訓(xùn)練權(quán)重，采用初始學(xué)習(xí)率為0.005，Вatchsize 設(shè)置為16，使用的網(wǎng)絡(luò)優(yōu)化器為Adam 算法，最大迭代次數(shù)epoch 為200 次。

2.2.2 模型訓(xùn)練評估指標(biāo)

本文主要使用mAP 值來對交通標(biāo)志識別模型進(jìn)行評估。平均精度（AP）是精確率和召回率曲線下的面積，是用于評估模型的性能指標(biāo)之一。mAP 則為AP 的平均值，主要用于衡量每個交通標(biāo)志類別識別結(jié)果的AP 值。一般情況下，模型的mAP 值越高，說明模型的性能越好。實際平均精度的計算公式如下所示：

其中，k 代表圖片數(shù)目，p(k)表示識別出k 個圖片是準(zhǔn)確率的值，Δr(k)表示識別的圖片由k-1 增加到k 時召回率的變化情況。

（2）識別速度評估

本文針對識別速度的評估主要是通過Python 代碼內(nèi)的time 庫來實現(xiàn)計時功能。根據(jù)代碼記錄的交通標(biāo)志識別前后的時間間隔來得到識別每張圖片的交通標(biāo)志所需要的時間。

■2.3 模型訓(xùn)練結(jié)果與分析

通過使用TT100K 數(shù)據(jù)集對模型進(jìn)行訓(xùn)練，可以得到較好的交通標(biāo)志識別精度，具體的模型訓(xùn)練結(jié)果如圖6 所示。

圖6 模型訓(xùn)練結(jié)果圖

從 圖6 可 知，以STN-GhostNet 為主干網(wǎng)絡(luò) 的YOLOv5s 模型對交通標(biāo)志圖像具有較好的識別效果，識別的mAP 值最終可以在交并比為0.5 的情況下達(dá)到98.4% 的mAP 值，而 在交并比為[0.5:0.95]時可以達(dá)到90.1%的mAP 值，說明其具有較好的交通標(biāo)志識別效果。部分交通標(biāo)志圖像識別的結(jié)果如圖7 所示。

圖7 交通標(biāo)志識別結(jié)果圖

同時本文提出的交通標(biāo)志識別模型經(jīng)過模型輕量化中的剪枝操作后（剪枝率為80%），模型尺寸從原先的14.355KВ 減少為了6.331KВ，共減少了約55%的模型尺寸；且經(jīng)過測試，該模型進(jìn)行交通標(biāo)志識別的速度約為23ms/幀，相較于原YOLOv5s 的識別速度有了較大幅度的提升，有助于其后續(xù)實際應(yīng)用于自動駕駛系統(tǒng)中。

3 結(jié)語

本文在原YOLOv5s 模型的基礎(chǔ)上，針對實際交通標(biāo)志識別的需求，結(jié)合STN 網(wǎng)絡(luò)與GhostNet 網(wǎng)絡(luò)的優(yōu)點，提出了一種改進(jìn)型的交通標(biāo)志識別網(wǎng)絡(luò)STN-GhostNet。并令其作為YOLOv5s 的主干網(wǎng)絡(luò)，利用輕量化的技術(shù)，對模型的尺寸與參數(shù)進(jìn)行修剪，使模型最終可以更好地滿足識別交通標(biāo)志時效性與準(zhǔn)確性的需求，同時其輕型的特點更適合部署于實際的自動駕駛系統(tǒng)中。目前隨著YOLO 系列的逐步發(fā)展，其識別性能逐漸提高，針對YOLO 的邊緣端移植也有了初步的實現(xiàn)，但僅限于較低的YOLO 版本，后續(xù)該模型在自動駕駛方面的實際應(yīng)用依舊任重道遠(yuǎn)。