基于Transformer的陶瓷軸承表面缺陷檢測方法*

安 冬,胡榮華,王麗艷,邵 萌,李新然,劉則通

(沈陽建筑大學機械工程學院,沈陽 110168)

0 引言

Si3N4軸承滾子具有硬度高、熱膨脹系數低、自潤滑等優良性能。廣泛應用于軍事、航空航天等領域的關鍵動力設備[1]。Si3N4軸承滾子易碎,在生產過程中容易產生缺陷,如劃痕和凹坑。隨著基于人工智能技術的缺陷檢測系統的大規模應用[2],深度學習方法在缺陷檢測任務上取得了較為高效準確的檢測精度[3]。但由于軸承滾子表面曲率與檢測設備限制,當軸承滾珠尺寸較小時往往會導致表面缺陷特征模糊不清,從而影響檢測精度。因此,從低分辨率圖像恢復高分辨率表面圖像,對提高缺陷檢測精度具有重要意義。

圖像超分辨率是旨在從單幅低分辨率(low resolution,LR)輸入圖像生成高分辨率(high resolution,HR)圖像的基本任務[4]。最近的研究中,基于卷積神經網絡網絡(convolutional neural network,CNN)的超分辨率方法的效果尤其突出[5]。但CNN單純使用滑動卷積窗口作為特征提取方式,單一窗口的感受野有限,只注重提取局部信息,而忽略特征之間的關系,既全局信息[6]。在實際的工業生產環境中,由于缺陷樣本的出現伴隨著隨機性,且同一缺陷類型的缺陷尺度大小并不一致。CNN有限的感受野將不可避免地影響全局特征結構的恢復,導致重建特征產生模糊效果[7]。

DOSOVITSKIY等[8]開創性地提出了一種以編碼器和解碼器為主要網絡結構的深度學習模型,并將其命名為ViT(Vision Transformer)。YANG等[9]最早將Transformer結構引入超分辨率任務中,提出了一種針對圖像紋理恢復的超分辨率方法。LU等[10]發現,相比于CNN網絡,ViT擁有更加優秀的上下文信息建模能力,并提出了一種可以捕獲長距離上下文相關性的神經網絡。

本文提出了一種結合全局自注意力模塊與殘差學習函數的殘差ViT超分辨率模型。具體來說,受殘差學習思想的啟發,本文提出了一種集合CBAM模塊[11]的局部自注意力模塊改善Transformer對邊緣特征的提取能力。同時使用殘差學習策略改善網絡模型對邊緣特征的重建精度,并改進了殘差學習函數以提高重建精度。

1 相關理論

1.1 殘差學習

LRi=I(D(HRi))∈H×W

(1)

式中:D(·)是下采樣算子,I(·)是具有相同比例因子的相應插值算子。

(2)

在最近的研究,殘差學習方法相比于其他方法表現出了更好的優化特性,并且可以通過端到端的信息傳遞大大增加在深度網絡中的準確性。如式(3)～式(5)所示,對于雙三次插值方法,從插值的LR圖像到圖像的映射接近于一個相同的映射,特別是當D(·)和I(·)的比例系數較小時。因此,本文所提出模型,設計為學習殘差圖像R,預測圖像為SR。

R=HR-LR

(3)

(4)

(5)

1.2 Transformer

Transformer是一種新型神經網絡結構,通過若干個帶有自注意力模塊的編碼器與解碼器的堆疊,神經網絡可以像人類的視覺系統一樣,在眾多信息中將注意力集中在重要的目標上,將權重分配給關鍵信息,而忽略其他不重要的信息。ViT(vision transformer)首次將Transformer用于圖像領域的網絡模型。它僅使用了Transformer的編碼器部分,通過將一系列經過位置編碼后的圖像子塊作為網絡輸入,實現了圖像分類任務。

ViT的自注意力機制使得它能夠準確地捕獲像素之間的長期相關性。且全局計算特征使得它能夠很好地對遠程上下文信息進行建模。雖然ViT在提取圖像的全局表示方面有很大優勢,但僅依賴圖像級的自注意機制仍然會導致局部細粒度細節的丟失。因此,如何有效地結合圖像的全局信息和局部特征,對于高質量的圖像重建非常重要[12]。

2 改進超分辨率ViT網絡

本文將ViT網絡引入圖像超分辨率任務中,提出了一種超分辨率ViT網絡,在傳統ViT網絡的基礎上,加強網絡表面細節特征提取能力,設計了一種增強局部特征注意力的局部自注意力模塊,網絡結構如圖1所示。

圖1 超分辨率ViT網絡結構

2.1 基于CBAM的特征提取模塊

CBAM是一個輕量級通用模塊,如圖2所示。CBAM模塊由通道注意模塊(channel attention module,CAM)和空間注意模塊(spatial attention module,SAM)組成,可以分別學習要注意的語義信息和位置信息,使神經網絡更加關注目標區域,抑制無關信息,提高缺陷邊緣的重建精度。首先,對給定中間特征圖m的輸入,通過通道注意力模塊獲得通道注意力圖MC(m)。然后將其與m相乘以獲得通道細化特征m′。同樣,以m′為輸入,將空間注意力模塊輸出的空間注意力圖MS(m′)乘以m′,得到最終細化輸出m″。整個過程可以總結為:

圖2 CBAM注意力模塊

m′=MC(m)?m

(6)

m″=MC(m′)?m′

(7)

2.2 L2多頭自注意力模塊

在深度學習中,利普希茨連續條件[13]作為常用神經網絡的約束,以控制網絡輸出相對于輸入的變化程度。利普希茨約束可以賦予模型對抗擾動的穩健性,并保證泛化邊界。利普希茨約束可以用于穩定訓練,例如譜歸一化。利普希茨常數L由下式給出:

(8)

然而,KIM等[14]的研究中發現,使用傳統點積多頭自注意力模塊的利普希茨常數可以是無界的,從而導致網絡梯度爆炸使訓練失效。為改進這個問題,本文提出了一種L2多頭自注意力模塊(LMA),如式(9)所示。如式(10)所示,將點積運算替換為歐幾里德距離,通過計算投影矩陣WQ和WK之間的矢量化L2距離消除梯度爆炸:

LMA(X)=concat[Attentionh(X)]W+b

(9)

(10)

最終本文提出的超分辨率編碼器結構如圖3所示。網絡結構如式(11)～式(14),對于輸入二維圖像x∈H×W×C在輸入時被分割為子塊序列,單個圖像子塊記為xp∈N×(P2·C),其中N=(H×W)/p2為序列長度。

(11)

圖3 局部多頭自注意力模塊

圖4 Si3N4軸承滾動體表面缺陷

(12)

(13)

y=LN(zn)=[y1,…,yN]

(14)

式中:LMA(·)為本文提出的L2自注意力模塊,LN為LayerNorm模塊。

在超分辨率ViT網絡中,本文使用逐像素MSE與MAE聯合損失作為損失函數。MSE定義為誤差平方的平均數,因此當誤差大于1時,MSE損失會更加敏感,網絡權重將以較快速度收斂。然而由于MSE中平方項的過度懲罰,MSE優化問題的解決方案通常缺乏高頻內容,這會導致在生成圖片中產生過于平滑的紋理。

為解決此問題,本文提出一種MSE與MAE聯合損失函數。首先,通過具有較大學習率的MSE損失函數訓練網絡,然后通過具有較小學習率的 MAE損失函數對網絡進行微調,這樣可以在加速訓練的同時略微提高準確率,聯合損失函數標識為:

(15)

(16)

3 實驗結果及分析

3.1 數據集建立

本節主要介紹了自制的Si3N4陶瓷滾珠表面缺陷數據集,并在此基礎上進行了表面圖像超分辨率實驗。該數據集是由60個失效6204陶瓷滾珠組成。包含4種常見的表面缺陷類型:劃痕、凹坑、磨損、雪花。使用Keyence VHX-1000超景深顯微鏡采集,如圖5所示。數據集共1312張顯微圖像,包含1096張缺陷圖像,分辨率為1600×1200。

圖5 超景深顯微鏡

3.2 評價指標

本文通過客觀評價與主觀評價對比本文所提出方法與其余方法的重建精度?？陀^評價指標方面,本文通過廣泛使用的峰值信噪比(PSNR)和結構相似度(SSIM)對重建圖像進行客觀評價。PSNR值越高,代表網絡重建精度約高。PSNR如式(17)所示。SSIM從包括亮度、對比度和結構等多個指標衡量圖像的相似度,SSIM值越接近1,說明兩者之間的相似度越大。SSIM如式(18)所示。

(17)

式中:MAX(SR)為圖像SR的最大像素值,一般取255。MSE(HR,SR)為真實圖像HR與超分辨率圖像SR的均方誤差。

(18)

同時,相對平均光譜誤差(RASE)[15]和視覺信息保真度(VIF)[16]也被用于對超分辨率ViT網絡的重建精度進行定量分析。相對平均光譜誤差(RASE)作為圖像融合領域中常用的評價指標。它表征了該方法在所考慮的光譜帶中的平均性能。VIF通過計算重建圖像和參考圖像之間的相互信息來預測主觀圖像質量。

3.3 實驗設置

本節將介紹超分辨率ViT的訓練策略與細節。在訓練過程中,在每個時期隨機裁剪16個大小為32×32的灰度圖像塊作為輸入,超分辨縮放因子設置為×4。網絡編碼器數量設置為4,所有Transformer塊的輸入和輸出特征維度為1024,MLP隱藏層維度為4096。模型使用Adam優化器訓練,動量等于0.9。

本文采用殘差學習策略,并對模型進行兩階段訓練,總迭代次數設置為2×105次。前90%次迭代使用MSE損失函數訓練網絡,初始學習率設置為2×10-4,每5×104次迭代后學習率減半。后10%次迭代使用MAE損失函數微調網絡權重,初始學習率設置為1×10-4,每2×103次迭代后學習率減半。訓練過程中使用平移、旋轉、拼接、縮放等數據增強手段,應用概率為0.8。

3.4 實驗結果分析

在實驗中,本文使用自建陶瓷球表面缺陷數據集中的10張圖像進行測試。超分辨率ViT與近年來幾種先進的超分辨率方法進行了比較。包括MESRGAN、SRCNN、SRGAN、VDSR。

定量分析實驗結果如表1所示,顯然超分辨率ViT在所有評價指標下都取得了最好的結果。在SSIM與PSNR上,與性能最接近的VDSR相比,本文方法的PSNR提高了1.11 dB,SSIM提高了0.008 7。這主要是由于本文在訓練前段使用MSE損失函數訓練網絡,因此在峰值信噪比上得到了較為顯著的提升。在VIF和RASE指標上,與Bicubic相比,本文方法的VIF提高0.029 4,RASE降低了0.023%。

表1 自建數據集上各算法的評價指標對比

如圖6所示,實驗還設置了超分辨率ViT與其他超分辨率模型在×4縮放因子下的直觀視覺比較。如圖所示,傳統重建方法無法兼顧紋理信息與邊緣信息。如SRGAN較好地處理了缺陷的紋理細節,但卻錯誤地平滑了梯度變化明顯的缺陷邊緣,從而導致邊緣模糊。相反,SRCNN對于高頻信息更加敏感,重建圖像獲得了更加銳利的邊緣特征,紋理細節卻過于平滑,導致圖像信息丟失。顯然,本文提出的超分辨率ViT網絡重建的圖像包含更準確的紋理細節。并且在在邊緣和線條特征中,本文提出的方法也獲得了與表現最好方法相當的重建性能。這得益于增加的CBAM特征提取模塊增強了對局部細節的感知,使模型可以從特定的區域學習更多信息。

圖6 在自建表面缺陷數據集上的對比實驗

4 結論

為了實時、準確地檢測氮化硅陶瓷滾動體的表面缺陷,本文提出了一種基于Transformer的超分辨率重建網絡。主要方法包括以下幾點:①針對Transformer對圖像細節特征提取能力差的問題,在網絡中插入混合域CBAM注意力模塊,提高網絡對圖像紋理與細粒度特征提取能力。②改進傳統ViT模型中的自注意力模塊,提出了一種改進L2自注意力模塊。③為獲得真實可靠的實驗數據集,本文使用超景深顯微鏡構建了氮化硅陶瓷球表面缺陷數據集。實驗結果證明,本文提出算法在峰值信噪比與視覺信息保真度上均獲得較大提升,有效解決了表面缺陷邊緣模糊,提高后續檢測精度。