閔永智，李嘉峰，王 果

(1.蘭州交通大學(xué) 自動化與電氣工程學(xué)院，甘肅 蘭州 730070；2.蘭州交通大學(xué) 甘肅省人工智能與圖形圖像處理工程研究中心，甘肅 蘭州 730070)

由于列車運行中輪轂與鋼軌長期的接觸摩擦、沖擊,鋼軌缺陷多由表面開始并逐漸加劇。列車長期在有缺陷的軌道上運行,輪軌沖擊振動會引起強度、疲勞破壞,影響輪軌系統(tǒng)安全性和壽命,因此及時檢測軌面缺陷具有現(xiàn)實意義。

在軌面缺陷檢測中,常規(guī)的無損檢測方法有射線檢測、漏磁檢測、渦流檢測和超聲檢測等。射線檢測分析透過鋼軌的射線差異來判斷缺陷是否存在,可對缺陷進行準(zhǔn)確的定性定量,但由于設(shè)備結(jié)構(gòu)復(fù)雜,主要應(yīng)用于鋼軌生產(chǎn)加工時的前期缺陷檢測[1]。漏磁檢測通過獲取鋼軌淺表層缺陷處產(chǎn)生的漏磁場來檢測缺陷,可靠性高,但難以檢測鋼軌內(nèi)部缺陷。渦流檢測利用電磁感應(yīng)原理,當(dāng)鋼軌存在缺陷時感生渦流的強度和分布將改變,通過渦流變化確定缺陷位置[2],檢測速度快,無需耦合介質(zhì),但由于渦流的趨膚效應(yīng),只能檢測近表面的缺陷。超聲檢測通過超聲波在鋼軌中傳播的散射能量及返回能量的幅度來確定缺陷位置[3],檢測厚度大、靈敏度高,在我國鋼軌探傷中應(yīng)用廣泛,但受軌面幾何形狀、粗糙度的影響,對近表面的微裂紋檢測效果不佳。

機器視覺檢測法是一種自動化程度高的新興檢測方法,通過對鋼軌二維圖像進行預(yù)處理,提取形狀特征,將缺陷從圖像中定位或分割出來[4],具有非接觸、連續(xù)長時間工作于鐵路惡劣環(huán)境的優(yōu)點,可滿足大規(guī)模線路的檢測需求。賀振東等[5]對鋼軌圖像進行反向P-M擴散并與原圖差分,根據(jù)缺陷邊緣特性和面積濾波分割出缺陷圖像,可以很好地識別塊狀缺陷和線狀缺陷;陳后金等[6]根據(jù)鋼軌圖像的灰度梯度特征提取鋼軌缺陷的內(nèi)邊緣,通過膨脹運算定位缺陷部位,可有效區(qū)分正常軌面、缺陷區(qū)域和干擾區(qū)域;張輝等[7]使用基于曲率濾波和改進高斯混合模型的方法檢測軌面缺陷,提高了檢測精度;謝敏杰等[8]設(shè)計一種基于線陣CCD相機的軌面缺陷檢測系統(tǒng),用改進的最大熵閾值分割法分割鋼軌圖像,根據(jù)形態(tài)學(xué)上連通區(qū)域大小是否滿足要求檢測缺陷,獲得了較低的誤檢率和漏檢率。以上方法實現(xiàn)了特定場景下軌面缺陷的識別,但都需要人工設(shè)計和提取特征,算法復(fù)雜。

當(dāng)前,深度學(xué)習(xí)已在表面缺陷檢測上有所應(yīng)用。深度學(xué)習(xí)檢測法利用神經(jīng)網(wǎng)絡(luò)提取特征,通過大量樣本學(xué)習(xí)到更深層抽象的特征,檢測速度快,可實現(xiàn)端到端訓(xùn)練,無需復(fù)雜調(diào)參,但訓(xùn)練樣本不足時容易產(chǎn)生過擬合,影響模型對新數(shù)據(jù)的接受度。生成對抗網(wǎng)絡(luò)(Generative Adversarial Network, GAN)[9]為樣本不足提供了一種新的數(shù)據(jù)擴充方法,但其一直存在梯度消失、訓(xùn)練不穩(wěn)定等問題;Mirza等[10]提出CGAN,將條件變量引入其中指導(dǎo)生成過程;Arjovsky等[11]提出WGAN,用Wasserstein距離替換GAN損失函數(shù)中的J-S散度,解決了梯度消失問題;Radford等[12]將CNN與GAN相結(jié)合,提出深度卷積生成對抗網(wǎng)絡(luò)(DCGAN),通過限制網(wǎng)絡(luò)結(jié)構(gòu)很好地穩(wěn)定了訓(xùn)練過程,目前成為GAN研究的標(biāo)準(zhǔn)架構(gòu)。張曼等[13]針對遙感圖像少、標(biāo)注困難的問題,在DCGAN的判別器中加入多特征圖特征融合技術(shù),融合3種尺度特征來提高圖像質(zhì)量,再結(jié)合掩膜圖像與遙感背景圖像生成樣本標(biāo)簽;戚銀城等[14]在DCGAN中加入自我注意力機制使卷積提取樣本的模式更充分,在其損失函數(shù)中加入相對均值鑒別器和梯度懲罰來平衡生成器和判別器的能力,提高了生成的缺陷螺栓圖像的IS值;龍程[15]將殘差塊引入DCGAN減少棋盤效應(yīng);于文家等[16]在條件生成對抗網(wǎng)絡(luò)(CGAN)中引入自我注意力機制,不僅增加了圖像各個部分之間的連接,還可以生成類別標(biāo)簽;李秋麗等[17]在DCGAN的生成器中加入自我注意力機制使網(wǎng)絡(luò)有目的地學(xué)習(xí),在判別器中加入頻譜規(guī)范化來穩(wěn)定判別器的訓(xùn)練,提高了生成圖像的IS值;劉學(xué)平等[18]將自我注意力機制引入DCGAN的多個卷積層上進行文字圖像修復(fù),當(dāng)文字圖像中含噪聲、遮擋和筆畫丟失時仍能準(zhǔn)確修復(fù)圖像。

由于維護及時,從鐵路工務(wù)部門獲取的軌檢圖像中缺陷樣本較少,這使得深度學(xué)習(xí)方法在實際檢測中表現(xiàn)不佳。為此,羅暉等[19]采用裁剪、縮放和GAN方式對鋼軌數(shù)據(jù)集RSDDs進行擴充,并改進Faster R-CNN實現(xiàn)了多尺度檢測與識別;呂邦歡[20]通過添加噪聲、Gamma變換、旋轉(zhuǎn)變換和DCGAN擴充軌面缺陷圖像,再使用Faster R-CNN檢測軌面缺陷,但均存在生成缺陷圖像過于相似、多樣性不足的問題。

受上述方法啟發(fā),本文將自我注意力機制與通道注意力機制引入DCGAN,提出一種基于注意力機制的深度卷積生成對抗網(wǎng)絡(luò)(Attention-DCGAN)用于擴充鋼軌表面缺陷圖像。該網(wǎng)絡(luò)使用注意力機制捕獲圖像中長距離尺寸依賴性和通道依賴性,彌補了卷積只能處理局部鄰域信息的缺點,相比DCGAN,生成的缺陷圖像多樣性顯著增強。

1 相關(guān)理論

1.1 生成對抗網(wǎng)絡(luò)

GAN是一種生成式模型,主要由生成器G和判別器D兩部分組成,其工作原理見圖1。生成器接收一個隨機噪聲z,然后訓(xùn)練它去擬合原始樣本的分布,使其越來越接近真實樣本;判別器用來分辨生成樣本和真實樣本之間的差異。理想情況下,二者相互對抗并不斷優(yōu)化各自網(wǎng)絡(luò)參數(shù),最終會使判別器難以判斷收到的樣本來自真實樣本還是生成樣本,即達到納什平衡D(G(z))=0.5,此時就可以認(rèn)為生成器學(xué)習(xí)到了原始圖像的數(shù)據(jù)分布。

圖1 GAN的工作原理

GAN的目標(biāo)函數(shù)為

Ez～Pz(z)[lg(1-D(G(z)))]

(1)

式中:V(D,G)為生成樣本和真實樣本的差異度;x為真實樣本輸入;G(z)為生成樣本;E為數(shù)學(xué)期望;Pdata(x)為真實樣本分布;Pz(z)為噪聲分布如高斯分布;D(x)為判別器判斷真實樣本是否真實的概率;D(G(z))為判別器判斷生成樣本是否真實的概率。

判別器的能力越強,D(x)應(yīng)該越大,D(G(z))應(yīng)該越小,這時V(G,D)會變大,其目的是求最大值(Dmax);生成器希望自己生成的圖片越真實越好,即D(G(z))盡可能大,這時V(G,D)會變小,因此式(1)對于生成器來說是求最小值(Gmin)。優(yōu)化生成器和判別器的目標(biāo)函數(shù)分別為

(2)

Ez～Pz(z)[lg(1-D(G(z)))]

(3)

DCGAN的網(wǎng)絡(luò)結(jié)構(gòu)見圖2。

圖2 DCGAN的網(wǎng)絡(luò)結(jié)構(gòu)

1.2 注意力機制

注意力機制是機器學(xué)習(xí)中的一種仿生數(shù)據(jù)處理方法,機器通過模擬人類閱讀、聽說中的注意力行為來自動忽略低可能、低價值的信息。具體來說是根據(jù)任務(wù)目標(biāo)對關(guān)注的方向和加權(quán)模型進行調(diào)整,在隱藏層增加注意力機制的加權(quán),使不符合注意力模型的內(nèi)容弱化或者遺忘。注意力機制已在深度學(xué)習(xí)各個領(lǐng)域廣泛使用。Chorowski等[21]通過注意力機制提高了長時間語音輸入的魯棒性;Hou等[22]提出一種有效的注意力機制來實現(xiàn)在線語音識別;Vaswani等[23]提出一種單純基于注意力機制的網(wǎng)絡(luò),完全省去了遞歸和卷積,在翻譯任務(wù)中表現(xiàn)良好;Yao等[24]提出一種時間注意力機制,在視頻描述任務(wù)中自動選擇最相關(guān)的時間段;Zhang等[25]設(shè)計一個對稱的全卷積網(wǎng)絡(luò)用于目標(biāo)檢測中顯著性特征的提取。

注意力模型可以理解為一個查詢Query到一系列鍵值對Key-Value的映射,其原理見圖3。其中Source為輸入序列,由一系列Key-Value組成,同時網(wǎng)絡(luò)還存在一種查詢向量。注意力機制本質(zhì)上是通過相似函數(shù)計算Query和各個Keyi的相關(guān)性,得到每個Keyi對應(yīng)Valuei的權(quán)重,最終將權(quán)重和對應(yīng)的Valuei加權(quán)求和,即

圖3 注意力機制原理

Attention(Query,Source)=

(4)

式中:Lx為序列長度;Similarity為計算Query和每個Keyi的相似度的函數(shù)。

2 注意力深度卷積生成對抗網(wǎng)絡(luò)

2.1 自我注意力模塊

由于卷積核大小的限制,卷積只能處理局部鄰域信息,因此單獨使用卷積層難以有效構(gòu)建圖像的長距離依賴關(guān)系。自我注意力(Self-Attention)機制[26-27]是注意力機制的變體,它是一種非局部機制,可以更好地利用全局信息去生成圖像。將其引入DCGAN,可以使遠(yuǎn)距離的要素位置信息用于生成詳細(xì)信息,判別器也可以檢查圖像中距離較遠(yuǎn)的細(xì)節(jié)特征是否一致。計算機視覺中自我注意力機制的原理見圖4。

圖4 自我注意力機制的工作原理

首先通過1×1卷積對上一個隱藏層中的特征圖x∈RC×N進行線性變換和通道調(diào)整轉(zhuǎn)換為三個特征空間f、g、h。其中f、g用來計算關(guān)注度,將f的輸出轉(zhuǎn)置并和g的輸出相乘,經(jīng)過softmax歸一化得到注意力映射,并和線性變換后的原始特征h逐像素相乘,再通過1×1卷積調(diào)整通道數(shù)就得到自我注意力特征圖,其數(shù)學(xué)表達式為

(5)

[sij]=[f(xi)]T[g(xj)]

(6)

式中:[f(xi)]=Wfxi,[g(xi)]=Wgxi為從不同權(quán)重矩陣的圖像特征提取到的兩個特征空間;βj, i為合成第j個區(qū)域時,模型關(guān)注第i個位置的程度。

注意力層的輸出為

o=(o1,o2,…,oj,…,oN)∈RC×N

(7)

將自我注意力特征圖乘以一個比例參數(shù),并和原始卷積特征圖加權(quán)求和作為下一隱藏層的輸入

yi=γoi+xi

(8)

式中:γ為可學(xué)習(xí)的標(biāo)量,初值為0,然后逐漸學(xué)會為非局部特征分配更多權(quán)重。引入它可以讓網(wǎng)絡(luò)首先依賴于鄰域信息,之后再逐漸把權(quán)重分配到遠(yuǎn)距離信息上,從而先學(xué)習(xí)簡單任務(wù),再逐步增加任務(wù)的復(fù)雜性。

2.2 擠壓激勵模塊

擠壓激勵(Squeeze-and-Excitation)操作[28]是一種具備通道注意力機制的模型,通道注意力用來衡量不同通道捕捉到的不同特征間的重要程度。卷積通過在局部接受域內(nèi)融合空間和通道信息提取特征,擠壓激勵通過對卷積特征的通道間的相互依賴建模來提高網(wǎng)絡(luò)性能。擠壓建立通道間的依賴關(guān)系,激勵重新校準(zhǔn)特征,二者結(jié)合強調(diào)有用特征抑制無用特征,其工作原理見圖5。

圖5 擠壓激勵的工作原理

首先通過一個卷積層將隱藏層中的特征圖x映射為W×H×C的特征圖,擠壓操作對輸入的W×H×C特征圖進行全局平均池化得到1×1×C大小的特征圖,C為通道總數(shù) ,通過聚集空間維度上的特征映射生成通道描述符,通道描述符生成響應(yīng)全局分布的信道特征,使每個特征圖都能利用其他特征圖的上下文信息,從而具有全局感受野,其計算式為

(9)

式中:i為高H;j為寬W;c為通道數(shù);uc(i,j)為空間維度W×H×C中的第c個特征圖;Fsq為壓縮操作;zc為將uc進行壓縮后的結(jié)果。

激勵操作使用兩個全連接層做非線性變換來生成每個通道調(diào)制權(quán)重的集合,其計算式為

s=Fex(z,W)=σ(W2δ(W1z))

(10)

式中:W1為第一個全連接層的權(quán)重,其維度是(c/r)×c;W2為第二個全連接層的權(quán)重,其維度是c×(c/r),其中r為通道壓縮倍率;W1z為第一個全連接層;δ為ReLU函數(shù),然后與W2相乘為第二個全連接層;σ為Sigmoid函數(shù);W為每個特征通道生成的權(quán)重;Fex(z,W)為對W和擠壓后的結(jié)果z的激勵操作;s用來刻畫uc的權(quán)重。

最后將每個通道對應(yīng)的權(quán)重乘到對應(yīng)通道的輸入特征上,其計算式為

(11)

2.3 網(wǎng)絡(luò)結(jié)構(gòu)

Attention-DCGAN生成器和判別器網(wǎng)絡(luò)結(jié)構(gòu)見圖6、圖7。

圖6 Attention-DCGAN的生成器網(wǎng)絡(luò)結(jié)構(gòu)

圖7 Attention-DCGAN的判別器網(wǎng)絡(luò)結(jié)構(gòu)

生成器將100維的噪聲通過一個全連接層后進行維度轉(zhuǎn)換(reshape)成(4, 4, 1 024)的三維張量,再先后通過兩個卷積核為4×4,步幅為2的轉(zhuǎn)置卷積層輸出(8, 8, 512),(16, 16, 256)的張量,之后加入一次自我注意力操作,輸出仍然為(16, 16, 256)。對得到的結(jié)果做卷積核為4×4,步幅為2,通道數(shù)為128的轉(zhuǎn)置卷積后再加入一次自我注意力操作得到(32, 32, 128)的張量。再通過同樣的轉(zhuǎn)置卷積得到(64, 64, 128)的特征圖后加入一次擠壓激勵層SEBlock。之后重復(fù)兩次上述轉(zhuǎn)置卷積并加入SEBlock,僅分別改變轉(zhuǎn)置卷積的通道數(shù)為64和32得到(256, 256, 32)的特征圖。最終經(jīng)過一個卷積核為3×3,步幅為1,通道數(shù)為3的卷積層后輸出(256, 256, 3)張量,經(jīng)Tanh處理后即為一個生成樣本。其中每個轉(zhuǎn)置卷積后做batchnorm操作,再用線性門控單元(GLU)[29]做非線性化處理,對于得到的結(jié)果經(jīng)最近鄰插值法進行上采樣得到想要的圖片大小,GLU可以看作是處于激活狀態(tài)下的一種簡化的ReLU單元,在相同的訓(xùn)練時間下,可以獲得比ReLU更高的精度。

判別器對輸入其中的(256, 256, 3)樣本用一個卷積核為3×3,步幅為1,通道數(shù)為16的卷積層轉(zhuǎn)化為(256, 256, 16)的三維張量,再先后經(jīng)三個卷積核為4×4,步幅為2,通道數(shù)為分別為32、64、128的卷積層后加入一次自我注意力操作和SEBlock得到(32, 32, 128)張量。然后經(jīng)一個卷積核為4×4,步幅為2,通道數(shù)為256的卷積層后加入一次自我注意力操作和SEBlock得到(16, 16, 256)的特征圖后再經(jīng)一個卷積核為4×4,步幅為2,通道數(shù)為512的卷積層后加入一次SEBlock得到(8, 8, 512)張量,最后用一個全連接層將其展開輸出一維結(jié)果。

3 實驗

為評估本文方法,使用自制軌面缺陷數(shù)據(jù)集和Type-I RSDDs數(shù)據(jù)集[30]進行實驗并和DCGAN進行對比。自制數(shù)據(jù)集拍攝于某廠內(nèi)鐵路線,包含掉塊和擦傷的缺陷圖片349張,制作過程中從同一水平高度拍攝,由于背景中的扣件、軌枕嚴(yán)重影響網(wǎng)絡(luò)訓(xùn)練,統(tǒng)一處理為僅含軌面區(qū)域的300×130像素圖片,部分圖片見圖8,Type-I RSDDs數(shù)據(jù)集包含大中小三種尺度的掉塊缺陷圖像。本方法主要適用于生成掉塊和擦傷圖片。實驗設(shè)備為Intel(R) Core(TM) i9-10920X CPU@3.50 GHz,128 GB運行內(nèi)存(RAM),NVIDIA GeForce RTX 3090 GPU,PyTorch1.8.0框架,Python3.7編譯環(huán)境。

圖8 部分自制數(shù)據(jù)集圖片

實驗中使用Adam優(yōu)化器,其中β1=0.5,β2=0.999,設(shè)置學(xué)習(xí)率為0.000 2,batch_size為1,iterations為100 000。在自制數(shù)據(jù)集上生成的圖像見圖9、圖10,在Type-I RSDDs數(shù)據(jù)集(圖11)上生成的圖像見圖12、圖13。

圖9 DCGAN在自制數(shù)據(jù)集上生成的圖片

圖10 Attention-DCGAN在自制數(shù)據(jù)集上生成的圖片

圖11 部分Type-I RSDDs數(shù)據(jù)集圖片

圖12 DCGAN在Type-I RSDDs數(shù)據(jù)集上生成的圖片

圖13 Attention-DCGAN在Type-I RSDDs數(shù)據(jù)集上生成的圖片

本文使用IS(inception score)[31]和FID(Fréchet inception distance)[32]來評價生成圖像的質(zhì)量。IS衡量生成圖片的清晰度和多樣性,該評分越高表示生成圖像質(zhì)量越好,其公式為

IS=exp(Ex～G(z)DKL(p(y|x)‖p(y)))

(12)

式中:x為生成的圖像;y為生成圖像輸入Inception網(wǎng)絡(luò)產(chǎn)生的向量;p(y|x)為生成圖像x被分類成某一類的概率分布;p(y)為生成的所有圖片在全體類別上的邊緣分布;DKL為兩分布之間的K-L散度。

FID計算真實樣本和生成樣本在特征空間之間的距離。首先利用Inception網(wǎng)絡(luò)提取特征,然后使用高斯模型對特征空間進行建模,再求兩個特征之間的距離,它的值越小表示生成圖像質(zhì)量越好,其公式為

(13)

式中:μdata、μg分別為真實樣本和生成樣本高斯分布的均值;∑data、∑g分別為真實樣本和生成樣本高斯分布的協(xié)方差。

原圖與DCGAN、Attention-DCGAN在自制數(shù)據(jù)集和Type-I RSDDs數(shù)據(jù)集上的IS、FID評價得分見表1,結(jié)果表明Attention-DCGAN比DCGAN生成圖片的質(zhì)量更好。從主觀角度來看,Attention-DCGAN相比DCGAN主要提高了缺陷分布的多樣性。在小樣本訓(xùn)練的情況下,DCGAN生成的圖像中缺陷較為單一,而Attention-DCGAN生成的圖像中缺陷數(shù)量不定,并創(chuàng)造了一些原圖中沒有的形狀,多樣性明顯提高。

表1 自制數(shù)據(jù)集和Type-I RSDDs上的IS/FID評價得分

4 結(jié)論

本文結(jié)合DCGAN和自我注意力、通道注意力機制提出一種注意力深度卷積生成對抗網(wǎng)絡(luò)(Attention-DCGAN),通過注意力機制增強了圖像中的通道依賴關(guān)系和遠(yuǎn)程尺寸間的依賴關(guān)系,提高了生成軌面缺陷圖像的多樣性。使用深度學(xué)習(xí)進行檢測任務(wù)時往往存在缺陷樣本不足的問題,通過GAN進行數(shù)據(jù)擴充是常用手段,但少量的缺陷樣本卻難以生成豐富多樣的缺陷圖像,難以提高目標(biāo)檢測、語義分割算法的泛化能力,Attention-DCGAN為這種矛盾提供了一種解決方案,有助于深度學(xué)習(xí)在檢測上的應(yīng)用。實驗結(jié)果表明Attention-DCGAN的性能優(yōu)于DCGAN,可以用其大量擴充軌面缺陷數(shù)據(jù)。雖然本文設(shè)計的Attention-DCGAN提高了生成缺陷圖片的多樣性,但與真實圖像比起來還是存在一定差距,因此如何生成高質(zhì)量的高分辨率圖像還需進一步研究。