基于輕量化模型的鋼軌扣件缺陷檢測系統(tǒng)

張 元,孟建軍,2,3,呂德芳,祁文哲,胥如迅,2,3,,陳曉強,2,3,

(1.蘭州交通大學機電技術研究所,甘肅蘭州 730070;2.甘肅省物流與運輸裝備行業(yè)技術中心,甘肅蘭州 730070;3.甘肅省物流及運輸裝備信息化工程技術研究中心,甘肅蘭州 730070;4.蘭州交通大學機電工程學院,甘肅蘭州 730070)

0 引言

鋼軌是鐵路軌道的主要組成部件,鋼軌扣件作為具有剛性扣壓件的鋼軌與軌枕的緊固裝置零件,是保障鐵路運營安全的重要組成部分[1]。研究表明,鋼軌扣件通常因為加工工藝不達標[2]或是長時間受惡劣自然環(huán)境甚至是惡意的人為因素的影響會出現(xiàn)松動或損毀的狀況[3]。

鋼軌扣件檢測以人工檢修為主,其檢測結果嚴重依賴巡檢工人的技術熟練程度和環(huán)境因素,誤檢率較髙。隨著機器視覺技術的廣泛應用,出現(xiàn)大量目標檢測模型并被使用于軌道狀態(tài)檢測工作。J.Chen等[4]提出了基于YOLO的一個級聯(lián)三階段的缺陷扣件檢測網(wǎng)絡,實現(xiàn)了較高的檢測率;Y.Li等[5]將深度可分離卷積和特征金字塔結合到YOLO V3網(wǎng)絡中,以提高缺陷檢測的精度。YOLO是一種典型的單階段模型,YOLO V5模型大小僅為27 MB,檢測速度快,但基于COCO數(shù)據(jù)集[6]的 mAP (mean average precision,平均精度的均值)性能指標僅約36%;若使用YOLO V4,盡管mAP 性能指標接近43%,但是檢測速度有所降低,并且模型體積高達245 MB。因此,為保證較高的檢測精度和檢測速度,提出一種改進YOLO V5模型,設計基于輕量化模型的鋼軌扣件缺陷檢測系統(tǒng),通過攝像頭拍攝和模型計算,最后將檢測信息反饋到顯示屏,并對比搭載未改進的原生YOLO V5模型的系統(tǒng)檢測結果,驗證本文方法的有效性。

1 系統(tǒng)總體設計

圖1為系統(tǒng)的總體設計結構及硬件結構,系統(tǒng)硬件部分由嵌入式設備Jetson TX2(包括充當視覺傳感模塊的板載CSI攝像頭)、內(nèi)存卡、鍵盤、鼠標和便攜式顯示屏組成。圖2為Jetson TX2的結構簡圖,其中的Tegra X2為整個嵌入式設備的計算處理模塊,為檢測和訓練過程提供圖像處理的算力;Jetson TX2自身搭載的CSI攝像頭對攝像頭視場范圍內(nèi)的鋼軌進行掃描,并拍攝鋼軌扣件;便攜式顯示屏通過VGA電纜線連接到Jetson TX2上,用于檢測信息的顯示;內(nèi)存卡中存放著系統(tǒng)的軟件部分,將內(nèi)存卡插入SD卡槽并開啟Jetson TX2,用鍵盤和鼠標對系統(tǒng)進行交互操控。

圖1 系統(tǒng)總體設計及硬件結構

圖2 Jetson TX2結構簡圖

圖3為系統(tǒng)軟件結構示意圖。Jetson TX2開機后進入內(nèi)存卡中存放的Ubuntu 18.04操作系統(tǒng),在輕量化模型主目錄下,攝像頭調(diào)用程序操控CSI攝像頭對鋼軌扣件進行圖像自動采集;訓練程序用于檢測工作開始前調(diào)用輕量化檢測模型對扣件圖像數(shù)據(jù)集進行訓練并生成一個權重文件;檢測程序用于調(diào)用輕量化檢測模型并使用訓練得到的權重文件對CSI攝像頭實時拍攝的扣件圖片進行檢測。

圖3 系統(tǒng)軟件結構

圖4為系統(tǒng)進行扣件缺陷檢測工作的流程,同時運行2個進程,分別實現(xiàn)對扣件進行自動拍攝和自動檢測,提升效率。進程都是循環(huán)自動執(zhí)行,當檢測工作結束時只需關閉進程即可。

圖4 系統(tǒng)檢測鋼軌扣件缺陷的流程

2 輕量化模型設計

圖5為改進后的輕量化YOLO V5模型。在用于訓練數(shù)據(jù)集時,將鋼軌扣件圖片壓縮成特征圖輸入,經(jīng)過主干網(wǎng)絡backbone的卷積,扣件被分成S×S個網(wǎng)格,每個網(wǎng)格預測N個邊界框和每個邊界框的置信度和類別概率;而用于檢測時,將整張圖像輸入模型,經(jīng)過backbone的卷積后形成多維特征圖,在Neck部分的Concat層進行特征圖融合后,再經(jīng)過UnSample層上采樣提高圖像分辨率[7],最后經(jīng)過Detector部分的Conv2d卷積后輸出所有扣件位置、類別和精度。改進點在于:將原生YOLO V5的主干網(wǎng)絡更換為MobileNet V3,再向MobileNet V3的bneck層中引入CA協(xié)同注意力機制[7](coordinate attention)并修改激活函數(shù)為Mish,形成CAMish-MobileNet卷積網(wǎng)絡(如圖5所示),并以此網(wǎng)絡構成輕量化模型。

圖5 輕量化檢測模型及輕量化網(wǎng)絡結構

2.1 卷積網(wǎng)絡的更換

MobileNet系列卷積神經(jīng)網(wǎng)絡[8]相較于Large版本,MobileNet V3-Small的基本單元bneck數(shù)量更少,在整體網(wǎng)絡的參數(shù)量比MobileNet V2減少約20%,同時精度提高6.6%。MobileNet V3以深度可分離卷積為基本單元構建網(wǎng)絡,深度可分離卷積由深度卷積和逐點卷積組成。深度卷積時,特征圖的每個通道由一個卷積核進行卷積,卷積核數(shù)量等于通道數(shù),其表達式如式(1)所示,F為輸入特征圖,C為寬W和高H的卷積核,特征圖在第a通道上的(x,y)坐標經(jīng)過寬和高2個方向的卷積核后得到該通道上的卷積核權重元素坐標(i,j)并輸出為特征圖G。

(1)

逐點卷積首先將卷積核尺寸設置為 1×1并使用深度卷積來提取每個通道的特征,再使用逐點的形式對通道的特征進行關聯(lián)。深度卷積與逐點卷積累加的深度可分離卷積D的計算量如式(2)所示:

D=DC·DC·M·DF·DF+M·N·DF·DF

(2)

式中:DF、N分別為輸入特征圖F的邊長和通道數(shù);DC為深度卷積核的邊長;M為逐點卷積的通道數(shù)。

相較于普通卷積,深度可分離卷積減少的計算量如式(3)所示:

(3)

當采用3×3的卷積核即DC=3時,深度可分離卷積的計算量約為普通卷積的1/9,達到輕量化的目標。

2.2 激活函數(shù)的修改

YOLO V5原生的激活函數(shù)是ReLU 6,若使用這種激活函數(shù)可能會出現(xiàn)梯度消失或梯度爆炸的問題。將激活函數(shù)修改為Mish,可解決包括MobileNet V3自帶的Hard-Swish在內(nèi)的大部分激活函數(shù)都無法解決的梯度消失或梯度爆炸的問題。計算公式如式(4)所示,x∈R時,f(x)無邊界,避免了取值封頂而導致的飽和。這樣平滑的Mish激活函數(shù)允許信息更好地深入,可減少卷積網(wǎng)絡不穩(wěn)定帶來的不必要的計算量,在性能較低的嵌入式設備上運行更流暢。在原本參數(shù)量就較少的MobileNet V3卷積網(wǎng)絡中使用Mish激活函數(shù),最終檢測精度比使用Hard-Swish和ReLU分別提高了0.494%和1.671%[9],達成了“輕量化但不損失精度”的目標。

f(x)=x·tanh[ln(1+ex)]

(4)

2.3 協(xié)同注意力機制的融入

包括MobileNet V3在內(nèi)的眾多卷積網(wǎng)絡都采用了SE注意力機制,SE注意力機制僅考慮通道間的信息而忽略位置信息,致使全局信息全部壓縮到通道中,當圖片中出現(xiàn)多個扣件樣本時會出現(xiàn)預測框位置不準確的情況。

為避免上述情況,引入一種高效注意力機制CA[7],它對于尺寸為C×H×W的特征圖輸入按照坐標為(i,j)的元素的水平W方向和垂直H方向進行池化,分別生成尺寸為C×H×1和C×1×W的特征圖xC(h,i)和xC(j,w),如式(5)所示:

(5)

通過式(5)信息嵌入的變換后,接下來將C×H×1和C×1×W的特征圖zh和zw進行拼接,之后使用1×1的卷積模塊F1×1對其進行變換,將其維度降低為原來的C/r,經(jīng)σ(即Sigmoid函數(shù))激活得到特征圖f,如式(6)所示:

f=σ(F1×1([zh,zw]))

(6)

如式(7),此時利用另外2個代表空間方向的1×1卷積變換Fh和Fw分別將fh和fw經(jīng)σ激活得到2個方向的注意力權重gh和gw,輸入得到CA協(xié)同注意力機制的最終輸出yc如式(8)所示:

(7)

(8)

CA注意力機制同時考慮2個空間位置(i,j)方向的信息xc和yc,幫助模型更好地定位和識別目標,維度的降低帶來更小的計算開銷,使檢測模型輕量化。

3 系統(tǒng)部署與測試

3.1 系統(tǒng)部署

如圖6所示,按照圖1的硬件結構方案搭建鋼軌扣件缺陷檢測系統(tǒng)的硬件部分,圖6中細節(jié)之處在圖7中展示,其中Jetson的具體型號配置為NVIDIA Jetson TX2 8 GB搭載四核ARM Cortex-A57 MPCore處理器、配有256核心NVIDIA Pascal架構及8 GB 128位LPDDR4內(nèi)存,內(nèi)存卡型號配置為SanDisk Extreme PRO 32 GB。按照圖3的軟件結構方案構造鋼軌扣件缺陷檢測系統(tǒng)的軟件部分,將訓練程序train_model.py、調(diào)用攝像頭程序camera.py、檢測程序detect_defects.py置于輕量化YOLO V5檢測模型的主目錄下,如圖8所示。最后將系統(tǒng)安置在圖9的軌道巡檢小車上,在鐵路上駕駛時實現(xiàn)鋼軌扣件的自動拍攝與缺陷自動檢測。

圖6 系統(tǒng)硬件部分實物圖

圖7 Jetson TX2上的組件細節(jié)

3.2 模型訓練

數(shù)據(jù)集有1 376張不同工況的鋼軌扣件圖片,包括扣件正常無缺陷(標注為normal)360張、扣件螺母丟失(標注為loss_nut)322張、扣件彈條缺失(標注為loss_sp)374張、扣件彈條斷裂(標注為cracked_sp)320張,圖像分辨率從128 pixel×98 pixel至750 pixel×1 000 pixel。整個數(shù)據(jù)集的1 376張鋼軌扣件圖片共有2 032個標注框,各工況的標注數(shù)量如表1所示。

圖9 軌道巡檢小車

表1 不同工況扣件的標注數(shù)量 個

在訓練開始前,為了盡可能與未改進的YOLO V5清晰對比,超參數(shù)值設定使用與YOLO V5基本相同的配置[10]:模型初始學習率為0.005,訓練迭代100輪,Batch Size(批歸一化值)為8。

使用目標檢測領域中常用的mAP@.5和mAP@.5:.95[6]作為模型訓練的精度指標。

(9)

式中:n為類別數(shù)量;APk為第k類的平均精度;而mAP即所有k類平均精度的平均值。

實驗中需預先設定預測框與真實框交并比IoU(intersection over union)閾值,mAP@.5是指IoU閾值取0.5時的mAP值,mAP@.5:.95是指當IoU閾值以0.05為步長,從0.5取至0.95時的mAP值的平均數(shù)。

運行訓練程序train_model.py,使用設定好的超參數(shù)開始對鋼軌扣件數(shù)據(jù)集進行訓練。如圖10所示,訓練過程中,mAP@.5和mAP@.5:.95隨著迭代次數(shù)的增加而增加。

(a)mAP@.5與迭代次數(shù)的關系

(b)mAP@.5:.95與迭代次數(shù)的關系

同理,使用原生YOLO V5訓練數(shù)據(jù)集。將改進前后的模型的訓練結果進行對比,mAP性能指標、模型大小和顯存占用的情況如表2所示。

表2 模型改進前后對鋼軌扣件數(shù)據(jù)集的訓練性能比較

與原生YOLO V5模型對比可知,經(jīng)過改進的輕量化YOLO V5訓練精度mAP指標更高,模型體積更小且占用更少的顯存,滿足輕量化需求。

3.3 系統(tǒng)測試

巡檢工人駕駛軌道巡檢小車在鐵路上行駛,同時操作部署在巡檢小車上的鋼軌扣件缺陷檢測系統(tǒng)。如圖11所示,打開終端,運行攝像頭調(diào)用程序開啟圖像采集進程,攝像頭對準鋼軌扣件,每隔10 s拍攝1張新圖片,本次測試拍攝了69張鋼軌扣件圖片,其中包含100個扣件樣本,圖片存放在fastener_pic目錄下。

圖11 系統(tǒng)拍攝鋼軌扣件

將3.2節(jié)中經(jīng)過輕量化模型訓練數(shù)據(jù)集得到的權重yolov5-lightweight.pt如圖8中置于輕量化檢測模型的主目錄下。同時打開另一個終端,運行檢測程序,對fastener_pic目錄下拍攝到的鋼軌扣件圖片進行檢測,檢測運行過程如圖12所示。

圖12 系統(tǒng)進行扣件缺陷檢測

100個樣本的扣件缺陷檢測結果如表3所示。

表3 系統(tǒng)檢測扣件缺陷結果 個

經(jīng)系統(tǒng)檢測輸出后的部分扣件缺陷識別結果實物如圖13所示。

(a)正確檢測為“扣件正常”

(b)錯誤檢測為“彈條斷裂”

(c)正確檢測為“彈條丟失”

(d)檢測為“彈條丟失”但精度較低

(e)正確檢測為“螺母丟失”

(f)錯誤檢測為“扣件正?！?/p>

檢測完成后,會在終端顯示平均每張扣件圖片的檢測時間和FPS(frames per second,幀速率)。如表4所示,為使用改進前后YOLO V5得到的權重進行扣件缺陷檢測的速度和FPS對比。結合上述檢測結果,證明搭載輕量化模型的本系統(tǒng)具有實現(xiàn)檢測速度更快、檢測準確度更高的輕量化表現(xiàn)。

表4 改進前后的模型檢測速度和FPS對比

3.4 系統(tǒng)改進

經(jīng)過3.3節(jié)的測試,系統(tǒng)最終對于100個鋼軌扣件樣本的缺陷檢測得到87%的準確率。進一步研究系統(tǒng)參數(shù)發(fā)現(xiàn),在模型訓練時,將超參數(shù)中的初始學習率進行修改,會影響最終系統(tǒng)檢測的準確度。如圖14所示,同樣拍攝100個鋼軌扣件樣本,使用初始學習率設置為0.01時訓練模型得到的權重文件來運行系統(tǒng)檢測工作,誤檢個數(shù)僅為4個,準確度較學習率設置為0.005時提升10.3%。

圖14 誤檢個數(shù)與初始學習率的關系

4 結束語

設計輕量化模型,對YOLO V5目標檢測模型的主干網(wǎng)絡和激活函數(shù)進行修改,并融入?yún)f(xié)同注意力機制。調(diào)用經(jīng)過訓練扣件數(shù)據(jù)集得到的權重文件,對CSI攝像頭拍攝到的鋼軌扣件進行缺陷檢測。

經(jīng)實驗測試,使用所設計的鋼軌扣件缺陷檢測系統(tǒng),對單張扣件圖片檢測時間為56.8 ms,誤檢率為13%。對系統(tǒng)進行改進,增加初始學習率數(shù)值,發(fā)現(xiàn)將初始學習率設置為0.01進行訓練得到的模型用于系統(tǒng)中,誤檢率僅為4%,相比系統(tǒng)改進之前,檢測精度更高且速度不減,對算力和存儲空間等硬件條件要求極低。