一種基于知識蒸餾的軌道檢測輕量化模型

湯文亮,曾建楊,何文晶

(華東交通大學 信息工程學院, 南昌 330013)

0 引言

我國鐵路、地鐵里程巨大且逐年上升,對交通運輸、經濟發(fā)展起著舉足輕重的作用[1]。鐵軌是鐵路和地鐵的重要組成部分,在運行過程中,因為天氣、過重的運輸負載等原因,鐵軌會隨時間受到不同程度的損耗,包括幾何結構缺陷、鐵軌部件缺陷、鐵軌表面缺陷[2]。常見的鐵軌表面缺陷包括疤痕、裂紋、波紋擦傷、褶皺、剝落等[3]。如果出現(xiàn)此類鐵軌缺陷,要及時進行排查修繕,不然容易留下隱患,引發(fā)重大事故。

目前,鐵軌檢測多采用人工檢測,效率低、成本高,而且還存在一定的安全隱患。還有一些傳統(tǒng)的軌道缺陷檢測方法,如磁粉法、電渦流法、超聲檢測法等[4],這些方法能滿足部分的檢測要求,但是容易受到外界環(huán)境的干擾[5],準確率也不是很高。近些年,隨著人工智能的快速發(fā)展,特別是在深度學習和計算機視覺[6]領域,涌現(xiàn)出了很多識別速度快、準確率高的神經網絡模型,大量的學者基于這些模型做了應用和優(yōu)化[7-8]。計算機視覺已經逐漸應用于鐵軌缺陷檢測鄰域,相較于傳統(tǒng)方法,檢測精度有了很大提升。戴鵬等[9]提出了一種無砟軌道扣件快速定位算法和識別扣件缺陷的半監(jiān)督機器學習方法,扣件缺失、彈條折斷、彈條移位3類缺陷的檢出率達到95%。王子冠等[10]基于MaskR-CNN[11]模型,綜合運用Sobel算子和Hough變換識別圖像中的鐵軌,對軌道進行區(qū)域識別,但是存在識別速度較慢,邊緣存在毛刺,影響識別效果的問題。侯博文等[12]運用遷移學習的方法,基于殘差網絡模型,實現(xiàn)對鋼軌、扣件及支承塊3種軌道結構的病害識別。

雖然目前鐵軌缺陷檢測的人工智能算法已經取得了一定的成果,但大多模型都是通過增加神經網絡的層數(shù)和參數(shù)讓模型獲得更強的特征提取能力和學習能力,這樣導致模型的參數(shù)存在巨大的冗余。因此,對網絡模型進行壓縮是非常有必要的,特別是在對檢測時間有嚴格要求的鐵軌檢測鄰域。Hinton等[13]提出可以用一個大的教師模型去指導小的學生模型,從而達到壓縮模型的目的,稱為知識蒸餾。針對上述問題,本文中提出了一個可用于鐵軌檢測的小型神經網絡模型,并利用知識蒸餾的方法和SAM優(yōu)化算法去訓練該模型。該模型參數(shù)較小,可以在眾多移動終端部署,快速地對鐵軌進行缺陷檢測。

1 相關工作

1.1 數(shù)據(jù)處理

鐵軌缺陷數(shù)據(jù)集來自kaggle上的Railway Track Fault Detection數(shù)據(jù)集,其將鋼軌磨損、扣件脫落、鋼軌斷裂、移位等病害分為defective類,將完好無損、沒有缺陷的分為no defective類,如圖1所示。

圖1 軌道缺陷檢測數(shù)據(jù)集

數(shù)據(jù)集中雖然類別較少,但一個類別中包含幾種病害類型,如裂縫、鐵枕斷裂、道床磨損、路基變形、扣件脫落等在本數(shù)據(jù)集中都有表現(xiàn)。數(shù)據(jù)集信息跨度較大,且各個圖片的長寬比、像素大小也不一樣,這樣的數(shù)據(jù)集對于模型的特征提取能力和泛化能力有很高的要求。因此,在本數(shù)據(jù)集中,如何在不同病害且特征分布不一的圖像中找出它們的共性,并提取出影響結果的關鍵共同特征,對于一個神經網絡模型是巨大的挑戰(zhàn)。由于鐵路里程巨大,鐵路檢測的重點是要找出哪一段路線出現(xiàn)了問題,以便及時進行排除和維修,所以能夠快速準確地檢測出具體哪一段軌道存在安全隱患就顯得尤為重要。

原始數(shù)據(jù)集中分為訓練集、測試集、驗證集,每一類各有150、11、31張圖片,由于數(shù)據(jù)集較少,將它們放在一起進行處理。這些數(shù)據(jù)用來訓練現(xiàn)在的網絡模型容易產生過擬合的現(xiàn)象。Wong等[14]提出數(shù)據(jù)增強對于提高性能和減少過度擬合有更大的好處。因此,使用增加亮度、增加銳度、改變對比度、橫移、翻轉和GAN[15]等數(shù)據(jù)增強方法對數(shù)據(jù)集進行擴充。數(shù)據(jù)集從192分別擴充為原來的10倍,即1 920張圖片,再將這些圖片按8∶2的比例分為訓練集和測試集。所以有1 536張圖像作為訓練集,384張圖像作為測試集。由以上數(shù)據(jù)集可以得出,預處理過程中主要有3個難點：① 數(shù)據(jù)集較少,模型容易出現(xiàn)過擬合現(xiàn)象;② 圖像的尺寸和拍攝的角度多樣,關鍵特征的坐標位置難以確定,且圖像是在自然光下進行獲取的,因此出現(xiàn)很明顯的陰影,影響模型識別;③ 每類樣本中包含多類鐵軌病害,圖像特征之間有較大差別,如何找出這些圖像共同的細小特征是一個巨大的挑戰(zhàn)。

1.2 SAM最小化銳度優(yōu)化算法

對于現(xiàn)在的主流模型,訓練損失一般來說是非凸的,有多個局部甚至是全局的極小值,不同的極小值模型的泛化性能也有顯著的差異。主流模型大多參數(shù)巨大,損失函數(shù)在迭代訓練的過程中,對模型的泛化能力的保障是很小的。在實際研究中,如果只追求最小的損失值,很容易導致模型的質量不佳。Foret等[16]提出將損失景觀的幾何圖形和一般化聯(lián)系起來,引入銳度sharpness概念,在訓練模型時同時最小化損失函數(shù)和損失銳度。SAM函數(shù)可以高效、簡便地在具有均勻低損耗值的鄰域內尋找參數(shù),極大地提高了模型的泛化能力。

根據(jù)PAC定理,生成的訓練集從分布上看極大概率滿足式(1)：

(1)

式中：h是一個單調遞增函數(shù)。將不等式右邊展開為：

(3)

最后,將sharpness的梯度值進行迭代即可最小化銳度。

1.3 知識蒸餾方法

隨著大數(shù)據(jù)和深度學習的快速發(fā)展,神經網絡模型越來越趨向于深度化。已經有不少學者證明神經網絡越深越大,該模型就越容易訓練[17]。但是在實際部署到用戶設備和終端中時,設備可以存儲的空間是有限的,擁有的計算資源也是有限的,參數(shù)量巨大的模型很難部署到終端,而且模型計算速度也難以得到保障。模型的壓縮引起大家的關注,知識蒸餾是模型壓縮的一個分支。知識蒸餾可以通過教師模型來指導小型學生模型的方式,讓學生模型在教師模型提取的特征和先驗知識的基礎上進行迭代訓練。學生模型通過知識蒸餾可以很好地解決小型模型難以訓練的問題,一般性能可以達到甚至超過教師模型,而且網絡結構更加簡單,參數(shù)更少,模型的計算速度也更快。

知識蒸餾一般需要構建2個神經網絡模型,即教師模型和學生模型,教師模型是層數(shù)深、參數(shù)量大、結構復雜的神經網絡模型,學生模型是層數(shù)淺、參數(shù)量小、結構精簡的神經網絡模型。教師模型一般易訓練且表現(xiàn)效果好,而學生模型計算速度快且占用存儲資源小。

將數(shù)據(jù)集輸入到教師模型中迭代訓練,得到一個較好的結果,再將數(shù)據(jù)集輸入到訓練好的教師模型中,將最后一個全連接層輸出的向量除以溫度(T)后再作softmax計算,得到軟標簽值。將軟標簽值和真實標簽值進行對比,得到loss1;再將數(shù)據(jù)集輸入學生模型,將得到的預測值和真實值進行對比,得到loss2,將2部分的損失相加,得到知識蒸餾的總損失,對總損失進行梯度下降,調整參數(shù),訓練學生模型,流程如圖2所示。

圖2 知識蒸餾訓練流程

神經網絡模型在進行分類任務時,最后一層的神經元輸出信息都會經過一個softmax函數(shù),如式(4)所示,它將神經元的輸出直接映射成該圖像類別的概率分布。在知識蒸餾中,這樣經過softmax輸出的概率信息稱為硬標簽信息。

(4)

式中：Pi是教師網絡模型輸出該圖像為第i類的概率;xi是最后一層神經元輸出的信息。

由于softmax函數(shù)只輸出概率分布的獨熱編碼,如果直接用教師模型的softmax輸出指導學生模型訓練,顯然會損失教師模型在計算中得到的各類別概率的豐富信息。基于此,知識蒸餾對其進行修改,增加了溫度系數(shù)T,如式(5)所示,使得輸出一個“軟化”后更加平滑的概率分布向量。

(5)

式中：T為溫度系數(shù),T越大,產生的概率就越平滑;Pi是教師網絡模型輸出表示該圖像為第i類的概率;xi是最后一層神經元輸出的信息。

在學生模型訓練時用到的總的損失函數(shù)是KLDivLoss,即計算真實標簽、教師模型輸出的軟標簽和學生輸出硬標簽的KL散度。KL散度又叫相對熵,可以很好地衡量2個隨機分布之間的距離,當2個隨機分布的差別增大時,它們的相對熵也會增大,文中用到的具體操作如式(6)所示。

(6)

式中：p是真實標簽的處理值;q是教師模型加學生模型的輸出,如式(7)、式(8)所示

(7)

(8)

式中：α、β是線性系數(shù);x1是教師模型輸出;T是溫度系數(shù);x2是學生模型輸出;y是真實標簽值。

1.4 缺陷檢測輕量模型

該模型結構簡單,參數(shù)量小,稱為SconNet6,包含6個卷積層,用于提取圖片特征信息,3個全連接層把特征信息映射到具體的類別,具體模型結構和卷積層參數(shù)如圖3、表1所示。

表1 SconNet6各卷積層的詳細信息

圖3 SconNet6模型結構

彩色鐵軌圖片的數(shù)據(jù)經過處理,得到一個3×224×224的張量作為輸入。6層卷積層提取不同維度的特征,全連接層根據(jù)特征預測結果。在每個Conv中都包含4項操作：卷積層、批量歸一化處理、激活函數(shù)和池化層。

先對輸入進行卷積處理,卷積層可以提取圖像特征信息,經過卷積后,圖像的通道數(shù)和尺寸大小都會相應變化。之后對數(shù)據(jù)進行批量歸一化處理,可以將離散的數(shù)據(jù)變成一個方差為1的正態(tài)分布,有效加快模型訓練時的訓練收斂速度,同時防止梯度消失和梯度爆炸,具體操作如式(9)所示。

(9)

將數(shù)據(jù)輸入到激活函數(shù)中,激活函數(shù)可以加入非線性因素,提高神經網絡模型的表達能力。本模型中用到的激活函數(shù)是ReLU函數(shù),如式(10)所示。

(10)

最后,用最大池化的方法將n*n個網格中的最大值當作該大網格的值,可以有效提取關鍵特征和壓縮模型參數(shù)。

2 實驗

2.1 實驗設定

本實驗用到的操作系統(tǒng)是Ubuntu,顯卡是顯存為12 GB的 NVIDIA GeForce RTX 2080Ti,實驗中用的語言和主要庫版本為python3.8.5,用python語言實現(xiàn)運行,版本為 3.8.5、pyTorch 1.8.0+cu111。

2.2 模型訓練與結果分析

將之前準備好的訓練集導入到VggNet16、ResNet18、ResNet50、DenseNet121模型進行訓練,測試集用來來驗證模型的性能,以上模型當作對照實驗。在訓練過程中,將epochs設為100,初始學習率設為0.001且隨著訓練而慢慢減小,batch size設置為16。在該實驗中使用的損失函數(shù)是交叉熵損失函數(shù),優(yōu)化器是對Adagrad算法進行擴展的Adadelta優(yōu)化器。以上4個模型訓練的測試結果如圖4所示。

圖4 各模型實驗結果曲線

由圖4可知,4種經典的網絡都在本數(shù)據(jù)集中表現(xiàn)出不錯的識別準確度。除VggNet16外,其他模型都在訓練到40多個epochs時開始收斂,識別準確率達到了90%以上。

但是訓練時的average accuracy曲線的抖動變化較大,有可能是訓練的學習率過大、batch size過小,或者是在模型梯度下降時優(yōu)化算法并沒有很好的效果。為此,在上述模型訓練時加入SAM最小化銳度優(yōu)化算法,結果如圖5所示。

圖5 SAM算法優(yōu)化實驗結果曲線

由圖5可知,在加入SAM優(yōu)化算法后,4個模型相較于傳統(tǒng)的優(yōu)化算法,average accuracy曲線都變得更加平滑了,說明模型更加穩(wěn)定,泛化能力也更強,并且在所有模型上的識別準確率也有一些提高。尤其是在VggNet16模型的實驗結果上,average accuracy曲線變得平滑了很多,而且準確率提升了近10個百分點。SAM優(yōu)化算法在數(shù)據(jù)集中能夠很好地優(yōu)化梯度下降的過程,提高模型的穩(wěn)定性和準確率。所以在下文SconNet6網絡的訓練中將使用SAM優(yōu)化算法進行訓練。

接下來用上述已經訓練好的模型當作教師模型,開始對學生模型SconNet6進行訓練。先對SconNet6進行單獨的訓練,當作對比實驗,再對教師模型做知識蒸餾指導訓練,結果如圖6所示。在本實驗中,知識蒸餾的溫度系數(shù)(T)為5,全能因子(α)為0.7,其他設置和之前實驗一樣。

圖6 實驗對比結果曲線

如圖6所示,選取在實驗中表現(xiàn)最好的模型DenseNet121當作對比,在SconNet6單獨進行模型訓練時,DenseNet121表現(xiàn)較好,平均準確率高出5個百分點,但在SconNet6進行訓練時,在經過40多個epochs后模型開始收斂,識別平均準確率超過DenseNet121。

各模型的實驗結果如表2所示,在所有的模型中,除了VggNet16以外,所有的模型識別平均準確率都超過了95%,其中采用知識蒸餾方法和SAM算法訓練的SonNet6模型表現(xiàn)最好,且模型大小僅為0.738 M,僅為DenseNet網絡的1/4。

表2 模型的識別準確率和參數(shù)

3 結論

針對鐵軌缺陷檢測模型過大、參數(shù)冗余導致檢測時間長和難以部署到移動終端的問題,提出使用知識蒸餾的技術和SAM優(yōu)化算法,訓練一個小型的鐵軌缺陷檢測模型。經過實驗證明,通過知識蒸餾的SonNet6模型在鐵軌數(shù)據(jù)集中的準確率和各項性能都遠超過單獨訓練時的模型,分類的平均準確率和運算速度超過了其他模型,極小的模型參數(shù)也大幅度降低了對計算資源的需求,可以很好地部署在眾多存儲空間不大的終端上,有很好的應用價值。