淺層卷積神經網絡融合Transformer的金屬缺陷圖像識別方法

唐東林 楊 洲 程 衡 劉銘璇 周 立 丁 超

西南石油大學機電工程學院，成都，610500

0 引言

在石油、化工、機械等工業領域，設備表面的金屬缺陷會對工業產品的質量造成極大影響，因此工業上對金屬缺陷檢測的關注度逐年增加。由于檢測速度慢、人工成本高、視力限制等，人工檢測已不能滿足當今行業的需求[1]。模式識別技術的發展極大地提高了缺陷檢測技術能力[2]。模式識別技術能夠自動學習已知樣本特征并對未知樣本進行自動識別，已廣泛應用于缺陷檢測識別領域[3-6]。

目前模式識別的缺陷檢測方法中基于深度學習的卷積神經網絡(convolutional neural network，CNN)應用最為廣泛[7]，它具有高度自動化的特點，能自動提取缺陷圖像特征，實現端到端學習。文獻[8]提出了級聯自編碼器結合CNN的金屬表面缺陷識別模型，優化了池化造成的信息丟失問題，但存在結構復雜、參數優化困難的問題；文獻[9]通過改進傳統CNN加強了背景和缺陷之間的對比度，并利用決策樹評估模型進行調參優化，但其參數量巨大，容易陷入過擬合狀態；文獻[10]改進了ResNet50，添加了可變形卷積，優化了分類結構，提高了鋼材表面缺陷識別準確率，但網絡結構復雜且計算量大；文獻[11]設計了輕量級CNN用于缺陷識別，并結合隨機森林增強了模型特征選擇能力，但網絡太淺圖像感受野太小，無法充分學習圖像語義信息；文獻[12]構建了多尺度特征重構網絡以獲取不同尺度的鋼板缺陷信息，但網絡分支較多，增加了網絡訓練時間；文獻[13-14]基于卷積神經網絡對鋼管焊縫缺陷進行自動識別，但其網絡結構過于單一，不具備識別不同尺度缺陷的能力；文獻[15]使用數據生成技術解決了缺陷數據不平衡的問題，并設計了輕量級CNN金屬缺陷分類器，但網絡中頻繁使用池化操作，造成了圖像信息丟失問題。總體上，深層CNN參數量多，計算量大，訓練成本高，而淺層CNN無法充分學習到圖像信息。此外，CNN通過卷積核往往只能考慮到圖像局部信息而忽略全局信息。

最近，自然語言領域擅長學習序列全局信息的Transformer模型[16]受到廣泛關注。文獻[17]將Transformer用于圖像分類任務，成功學習到圖像全局信息，在超大數據集上性能超越了CNN模型；文獻[18]引入蒸餾損失，進一步提高了Transformer圖像分類性能；文獻[19-20]將Transformer繼續加深，并探究了加深Transformer時保持精度提升的方法。目前Transformer在大型數據集上的表現已超越CNN，但其訓練成本非常高；在中小規模數據集上，由于Transformer直接學習全局信息且模型參數量大，學習難度高，通常處于欠擬合狀態，性能較差。

在此基礎上，本文提出了一種適用于中小規模數據集的，利用淺層CNN捕捉圖像局部信息，利用Transformer捕捉圖像全局信息，并引入通道注意力模塊SE(squeeze and excitation)[21]捕捉重要特征通道的金屬缺陷識別方法。實驗表明，在中小規模數據集上，相比傳統CNN和Transformer模型，該方法能更快速準確地實現金屬缺陷識別。

1 本文方法

CNN-Transformer模型結構如圖1所示。由模型整體框架(圖1a)可知：通過4層CNN學習圖像局部信息及其位置信息；利用Transformer捕獲全局像素點之間的相關性以學習圖像全局信息，同時引入SE模塊對重要通道進行重點關注；最終利用一層CNN融合所有圖像語義信息，通過全局平局池化進行降維，由全連接層(fully connected layer)實現金屬缺陷的自動識別。

(a)模型整體框架

1.1 CNN特征提取層

為解決深層CNN模型復雜度高、在中小規模數據集上容易過擬合的問題，本文采用4層淺層CNN進行特征提取，結構簡單且特征提取效率高，具體參數如表1所示。其中，輸入圖像尺寸為C×H×W，C、H、W分別表示圖像的通道數、高、寬，本文數據集均為單通道灰度圖，輸入C=1。本文采用補零操作的主要目的是使越靠近圖像邊緣的區域卷積后的像素值越小，實現圖像位置信息的學習。

表1 卷積層具體參數

由于特征提取中池化下采樣會造成圖像信息丟失，本文對此進行改進，通過設置卷積核步距為2的方式進行下采樣以保留圖像信息。為防止過擬合、梯度爆炸和梯度消失現象，卷積后利用BN(batch normalization)層進行批標準化處理，并通過ReLU激活函數進行非線性增強。

1.2 SE通道注意力模塊

不同特征通道包含不同圖像語義信息，為重點關注重要特征通道并抑制不重要的特征通道，本文引入通道注意力模塊SE，如圖1b所示。首先進行全局平均池化(global average pooling)以獲取通道級全局信息：

(1)

其中，Xin∈RC×H×W，表示輸入特征圖；X′∈RC×1×1，表示池化后的特征圖。

池化后進行通道權重計算，并將權重乘以原特征通道輸出最終特征：

Xout=Xin⊙σ(FC2(ReLU(FC1(X′))))

(2)

其中，FC1(C，C/4)、FC2(C/4，C)為兩層全連接層；σ為Sigmoid函數；⊙表示elementwiseproduct乘法；Xout∈RC×H×W，為最終輸出特征。

為優化梯度傳遞，本文在輸出特征與原特征之間設置了殘差模塊。

1.3 Transformer層

由于CNN模型通過卷積核僅能捕捉到圖像局部相鄰信息，因此本文融合Transformer以學習圖像全局信息，使模型在關注局部信息的同時考慮圖像全局信息，降低了訓練成本并提高了缺陷識別準確率。

Transformer來源于自然語言處理領域，能夠捕捉任意長度語言序列之間的相關性。本文融合Transformer用于缺陷圖像識別領域，其實現過程如圖1c所示：將三維圖像C×H×W分塊并展平為與自然語言同維度的二維序列，以訓練自然語言序列的方式訓練二維圖像序列，實現圖像全局信息的學習。展平后序列維度為N×(P2C)，P為圖像塊大小，N=HW/P2為圖像塊個數。

傳統Transformer需要嵌入位置編碼以學習序列之間的位置關系，本文在CNN層已通過補零操作實現自動學習圖像位置信息，因此無需嵌入位置編碼，簡化了Transformer模型結構。為使數據分布更加穩定，對二維序列利用LayerNorm層進行層標準化處理，再將其輸入多頭自注意力層。

多頭注意力機制(multi-head attention)是學習圖像全局信息的核心，為便于理解，先引入自注意力機制(self-attention)，計算過程如下：

Q=XPWqK=XPWkV=XPWv

(3)

(4)

其中，矩陣Wq、Wk、Wv∈RP2C×P2C；二維圖像序列XP∈RN×P2C；Q、K、V分別為查詢矩陣、匹配矩陣、值矩陣；Softmax為歸一化指數函數；d值取圖像序列維度P2C。

為擴展模型專注圖像不同位置的能力，引出多頭注意力機制，如圖2所示。多頭注意力機制具有多個自注意力頭，將Q、K、V按注意力頭數進行劃分，分別完成自注意力計算，并融合所有注意力頭的信息，更全面地考慮不同序列之間的相關性。

圖2 多頭自注意力機制Fig.2 Multi-head attention

為減小模型參數量與復雜度，本文去除了多頭自注意力層之后的前饋全連接結構，并設置了殘差模塊優化梯度傳遞。

傳統Transformer主要通過在二維圖像序列中額外插入一個class token向量進行訓練并用以分類。本文對其進行改進，去除class token，將二維圖像序列轉換為三維圖像，直接輸入分類層進行分類，簡化了模型結構。

1.4 缺陷分類層

分類層用于實現金屬缺陷的自動識別。如圖1a所示，本文利用一層卷積核大小為1×1的卷積層融合此前所有信息，不改變特征圖大小，輸出通道為1024；利用全局平均池化將特征圖降維至1024×1×1后將其展平；最后全連接層FC(1024，classes)將維度映射到類別數classes實現缺陷分類。

2 實驗數據設置

2.1 數據集介紹

為驗證本文方法有效性，引入了公開數據集NEU-DET[22]，同時通過實驗自建鋼板缺陷超聲數據集來驗證模型通用性。

(1)鋼材缺陷數據集NEU-DET。該數據集是東北大學發布的表面缺陷數據庫，如圖3所示，包含熱軋鋼帶的6種表面缺陷：軋制氧化皮(Rs)、斑塊(Pa)、開裂(Cr)、點蝕表面(Ps)、內含物(In)和劃痕(Sc)。該數據庫由1800張灰度圖像組成，每一類缺陷包含300張分辨率為200 pixel×200 pixel的灰度圖像。

圖3 NEU-DET數據集部分缺陷樣本Fig.3 Partial defect samples of NEU-DET dataset

(2)鋼板缺陷超聲數據集。為模擬工業中不同的金屬表面缺陷，在厚度11 mm的Q235鋼板上加工了2 mm、5 mm、8 mm三種深度的缺陷，每種深度具有不同缺陷形狀，如圖4所示。實驗采用頻率2.5 MHz、底部直徑20 mm的超聲探頭，探頭通過無缺陷的鋼板背面發射超聲信號并接收反射信號，利用數字式超聲探傷儀、示波器來讀取并記錄超聲缺陷信號。

圖4 缺陷鋼板和信號檢測設備Fig.4 Steel plate with defects and signal detection equipment

實驗采集了每種深度200個共600個缺陷信號，每個缺陷信號為1×16 384的一維波形信號。為滿足圖像處理實驗要求，將1×16 384波形信號轉化為128 pixel×128 pixel二維灰度圖像，如圖5所示。

圖5 超聲部分波形信號和數據轉化圖Fig.5 Waveform signal and data conversion diagram of ultrasonic

為充分利用所有數據樣本以減小實驗誤差，本文進行10次獨立實驗，每次均從原數據集中重新隨機抽取70%樣本作為訓練集，其余30%樣本作為測試集，取10次實驗平均結果作為最終實驗結果。其中，NEU-DET數據集以缺陷類別為分類標準，超聲數據集以缺陷深度為分類標準。具體數據集劃分結果如表2和表3所示。

表2 NEU-DET數據集劃分

表3 超聲數據集劃分

為提高模型的魯棒性，針對本文采用的所有數據集，在每一輪訓練中，對訓練集進行標準化處理，并以0.5的概率分別進行旋轉、平移、水平翻轉、垂直翻轉或同時進行水平垂直翻轉，實現每一輪參與訓練的圖像均不完全相同，以提高模型泛化能力；在測試集中則只進行標準化處理。

2.2 實驗參數設置

本文實驗平臺為Pycharm，深度學習框架為Pytorch，CPU為Intel(R)Core(TM)i5-10400F CPU @ 2.90 GHz，GPU為NVIDIA GeForce RTX 2060(6 GB)。根據不同數據集的規模對訓練參數進行微調，具體如表4所示。

表4 模型參數設置

3 實驗結果與分析

實驗中訓練集負責進行參數學習以擬合模型，測試集用以對模型性能進行評估，不參與訓練與參數學習，因此以下實驗均采用測試集的實驗結果作為模型評價指標。

3.1 缺陷分類實驗

為驗證本文模型對缺陷的識別性能，引入中小規模數據集NEU-DET進行實驗，模型準確率δ和損失Δ變化如圖6所示。因為迭代初期學習率相對較大，模型參數更新幅度大，導致準確率和損失有一定波動，但總體準確率呈現上升趨勢，迭代40次后模型便接近收斂狀態。同時模型結構簡單、參數量少，模型性能穩定，未出現過擬合現象。通過10次獨立實驗取平均值最終準確率為97.8%。

圖6 NEU-DET數據集性能曲線Fig.6 NEU-DET dataset performance curve

通過混淆矩陣展示模型的具體識別情況，如圖7所示。橫坐標為真實類別，縱坐標為預測類別，真實類別與預測類別相同則識別正確，即對角線值。模型整體識別準確率較高，但由于內含物(In)、點蝕(Ps)、劃痕(Sc)三類缺陷部分樣本存在相似特征，導致模型出現了識別混淆。

圖7 NEU-DET數據集混淆矩陣Fig.7 Confusion matrix of NEU-DET dataset

為驗證模型通用性，本文引入實驗采集的超聲數據集進行實驗，其準確率和損失曲線如圖8所示，在更小規模的超聲數據集上本文模型更快更穩定地達到了收斂狀態，這是因為超聲數據集類別更少，缺陷特征更為簡單，進行10次獨立實驗后平均準確率達到了更高的98.8%。

圖8 超聲數據集性能曲線Fig.8 Ultrasonic dataset performance curve

具體分類結果如圖9所示，模型對三種深度缺陷識別準確率均較高，證明了本文模型在中小規模數據集上具有較好的通用性。

圖9 超聲數據集混淆矩陣Fig.9 Confusion matrix of ultrasonic dataset

3.2 類激活熱力圖可視化實驗

為了驗證本文模型捕捉缺陷特征的能力，在公開數據集NEU-DET上可視化圖像類激活熱力圖[23]，即輸入圖片的不同區域對模型進行缺陷識別的影響程度，顏色越紅的區域越能主導識別結果，如圖10、圖11所示。

圖10 部分簡單缺陷類激活熱力圖Fig.10 Grad-CAM of some simple defects

由圖10可以看出，對于特征明顯的簡單缺陷，主導本文模型進行識別決策的核心區域明顯位于缺陷區域。

如圖11所示，對于尺寸小、數量多，以及一些特征不明顯的缺陷，熱力圖仍在缺陷區域顏色最深，證明了本文方法能有效捕捉缺陷信息。但也存在個別微小缺陷被遺漏的情況，這可能是影響識別精度進一步提高的原因，值得后續深入研究進行改進。

圖11 部分復雜缺陷類激活熱力圖Fig.11 Grad-CAM of some complex defects

得益于Transformer捕獲圖像全局信息的能力，使熱力圖整體上分布廣泛，并不僅集中于局部缺陷區域，而是從整個圖像收集了更多信息，更符合人類感知結果。

3.3 消融實驗

為驗證Transformer和SE模塊對模型性能的影響，將淺層CNN模型、引入SE模塊的CNN模型、CNN-Transformer模型和本文提出的引入SE模塊的CNN-Transformer模型進行對比實驗，評價指標包括準確率、損失、浮點計算量以及參數量，具體如表5所示。可以看出，Transformer、SE模塊都能有效改進CNN模型，在利用CNN捕捉圖像局部信息的基礎上，通過Transformer捕獲圖像全局信息，能夠以較少的參數量、較低的訓練成本提取完整的圖像特征，最大程度提高模型性能。引入SE模塊能在圖像通道層面做進一步特征過濾，加強重要通道的特征表達而抑制不重要的通道特征，進一步提高模型性能，且幾乎不增加額外的參數量與計算量。

表5 消融實驗

本文方法相對于淺層CNN模型僅有少量參數量和計算量的增加，但卻帶來了巨大的性能提升，實現了最高的缺陷識別準確率，證明了本文方法在中小規模數據集中能充分發揮性能。

3.4 網絡模型綜合對比實驗

為驗證本文模型的綜合性能，與CNN模型VGG16[24]、GoogleNet[25]、ResNet34[26]、MobileNetv2[27]、ShuffleNetv2[28]以及Transformer模型ViT[13]進行對比實驗。

實驗曲線如圖12和圖13所示，在本文的兩個中小規模數據集上，由于Transformer模型ViT直接學習全局圖像信息，學習難度大導致收斂速度慢。同時因其對數據量的要求非常高，本文數據集無法有效擬合模型參數，導致其性能較差且迭代曲線波動較大。

圖12 NEU-DET數據集性能對比曲線Fig.12 Performance comparison curve of NEU-DET dataset

圖13 超聲數據集性能對比曲線Fig.13 Performance comparison curve of Ultrasonic dataset

各模型具體性能指標如表6所示，訓練時間為數據集訓練耗時，推理時間為單張圖片推理耗時，對最優結果進行加粗。可以看出，主流CNN模型的參數量和計算量都比較大，在中小規模數據集中，其性能受限容易過擬合，參數最多的VGG16性能表現較差，表明CNN模型并不會因為可學習參數量多而取得更好性能。

表6 模型綜合指標對比

本文方法在中小規模數據集上取得了更好的綜合性能，MobileNetv2和ResNet34的準確率雖然和本文方法相近，但是本文模型網絡更淺、結構更簡單，推理時間更短；ShuffleNetv2參數量、推理時間與本文方法相差不大，但本文方法在保持高效率的同時能夠學習圖像全局與局部信息的優勢是其所不具有的，因而在準確率上表現更優。

4 結論

金屬缺陷識別技術在工業領域具有重要研究意義。本文針對中小規模金屬缺陷數據集提出了一種淺層CNN融合Transformer并引入SE模塊的缺陷識別模型，模型利用淺層CNN捕捉圖像局部信息，與主流深層CNN模型相比具有參數量少和計算量小的優點，可有效地避免過擬合現象發生；利用Transformer多頭注意力機制捕獲圖像全局信息，解決了CNN模型無法學習圖像全局信息的問題，同時引入通道注意力SE模塊實現重點關注重要通道。本文引入了NEU-DET數據集和超聲數據集來驗證模型性能，試驗結果表明，本文方法在準確率、參數量、計算量、推理時間等重要評價指標上都表現出良好的性能，對未來實現工業上金屬缺陷的在線高效率、高精度識別具有現實意義。