基于鄰域自注意力的鋼鐵表面缺陷分類算法

鞏 克,陸春月,柴子凡

(中北大學機械工程學院,山西 太原 030051)

鋼鐵行業是經濟產業的重要組成部分,鋼鐵表面缺陷不僅損害外觀和內部品質,嚴重時甚至會危及安全。然而,隨著生產水平的提高和需求量的增大,人工檢測已不能滿足鋼鐵表面缺陷檢測的要求,因此采用智能化、自動化的方法解決缺陷檢測問題已經是大勢所趨。隨著計算機視覺技術的發展,基于卷積神經網絡的檢測方法得到了廣泛應用。但是鋼鐵缺陷圖像成像模糊、分辨率較低,會導致網絡學習到的特征出現信息丟失、特征模糊以及易混淆的問題。為此,本文提出基于鄰域自注意力的鋼鐵表面缺陷分類算法,通過鄰域自注意力模塊與多尺度特征融合進一步提高了鋼鐵缺陷檢測的正確率。

1 相關工作

1.1 鋼鐵缺陷分類

傳統的鋼鐵表面缺陷識別方法主要使用小波變換、雙閾值二值化和決策樹等方法來分析和檢測圖像,但是適用性有限。近幾年提出的一些方法往往使用卷積神經網絡進行檢測,例如Boikov等[1]使用合成數據來訓練視覺任務的方法,在鋼工件表面缺陷的分類和分割方面都取得了良好的效果。Hao等[2]使用基于生成對抗網絡和注意力機制的方法來識別缺陷。Li等[3]提出了一種混合網絡架構(CNN-T),該架構合并了卷積神經網絡(CNN)和Transformer編碼器,在NEU-CLS數據集上取得了顯著的分類效果。

盡管上述方法在鋼鐵表面缺陷分類方面取得了不錯的效果,但這些方法所針對的鋼鐵缺陷數據集分辨率較高(通常為224×224)。當輸入圖像為低分辨率時,由于低分辨率圖像所包含的像素點數量較少,很多細節和信息都無法在圖像中表現出來,因此信息有缺失。同時,由于像素點的數量較少,低分辨率圖像的邊緣和輪廓通常比較模糊,這會使得一些細微的特征難以識別和區分,從而導致神經網絡所學到的特征模糊。

1.2 注意力機制

注意力機制的作用是讓系統學會從大量信息中把注意力放在感興趣或者高價值的地方,目前已經成功地應用于各種任務。例如,在2017年被應用于Transform模型中的自注意力機制[4],已成為大型模型發展的重要轉折點。此外,Hu等[5]提出了一種用于圖像分類的通道注意塊,以提高網絡的代表性。Wang等[6]則通過提出用于CNN的有效通道注意(ECA)模塊,成功實現了跨通道交互,增強了SENet的策略。為了建立通道注意和空間注意的雙重機制,Woo等[7]在SENet和ECANet的基礎上進一步加強了注意力模塊設計。在計算圖像自相關性時,往往使用鄰域相關性[8],但這種方法計算量大、模型復雜。而本文所提出的鄰域自注意力模塊通過簡單的連接實現上下文特征感知,省去了許多冗余的參數。

2 算法內容

圖1所示為網絡的總體架構,圖2所示為鄰域自注意力模塊的整體架構。鄰域自注意力網絡包括2個主要的可學習模塊:自相關計算模塊、上下文特征感知模塊。

圖1 基于鄰域自注意力網絡的總體架構

圖2 鄰域自注意力模塊架構

2.1 網絡總體架構

給定一組缺陷樣本圖像的情況下,使用卷積塊(Conv)來提取基本特征Z,Z∈H×W×C,其中H和W分別為特征的高度和寬度,C為通道維度。接著,引入鄰域自注意力模塊來增強基礎特征,并將增強后的特征與基礎特征進行殘差連接,得到自注意力特征A,A∈H×W×C。最后,將4個得到的特征進行多尺度特征融合,再經過輸出卷積層來恢復通道數,并通過全連接層進行分類。

2.2 自相關計算

為了獲取圖像中鄰域的自相似性,對基本特征Z,計算每個位置x處(x∈[1,H]×[1,W]及其鄰域中的值)C維向量的哈達瑪積,并將它們收集到自相關張量D中,D∈H×W×C。張量D可以表示為具有C1維向量輸出的函數:

(1)

式中:p∈[-dU,dU]×[-dV,dV],對應于滑動窗口中的相對位置,即2dU+1=U和2dV+1=V,其中dU和dV分別為垂直和水平方向上的最大位移,U和V為滑動窗口的高度和寬度。這里D并沒有保留U、V的維度,而把它看作通道特征的一部分,因此可以得到新的通道維度C1=U×V×C。

2.3 上下文特征感知

盡管自相關計算可以獲取圖像的自相似性,但它缺乏原始卷積特征所表示的局部語義線索。為了更好地捕捉語義對象的不同方面,對Z和D進行拼接,得到上下文語義特征G,G∈H×W×Cg,如下所示:

(2)

式中:G(i,j)為圖像中特定位置(i,j)的上下文語義特征,它包含了原始語義信息和自相似語義信息的綜合表達,能夠更好地反映圖像的語義特征。之后通過特征提取層來分析G中的上下文關系,并對提取出的特征張量進行再卷積操作,使用輸出層將特征通道數降至輸入通道數,得到更加緊湊的特征表示。上述兩個卷積塊的卷積核大小都為1×1,該卷積塊h(·)在沒有填充的情況下學習上下文關系,并聚集局部相關模式,從而將通道的維度恢復為C,使得輸出h(G)具有與Z相同的大小。將這兩種表示結合起來生成自注意力,表示為A∈H×W×C。

A=h(G)+Z

(3)

通過鄰域自注意力模塊對基本特征進行增強,有助于定位目標對象的重要區域并增強特征的可識別性。

2.4 多尺度特征融合

經過鄰域自注意力模塊可以得到4個加強的特征,它們分別代表不同尺度的自注意力特征。為了使特征圖信息完整并得到更好的嵌入特征,采用平均池化層和層歸一化進行特征處理,從而進一步提高特征的表達能力。

ai=LayerNorm(AvgPool(Ai))

(4)

式中:ai為處理后的特征,Ai為經過鄰域注意力處理后的加強特征。隨后,進行多尺度特征融合,將不同尺度的特征信息進行融合,以提高模型的分類準確性。最后,通過輸出卷積層恢復通道數,并后接全連接層進行分類。

a=concat(a1,a2,a3,a4)

(5)

y=FC(conv(a))

(6)

式中:a為融合后的特征,y為最終的分類結果,a1、a2、a3、a4為經過平均池化和層歸一化后的輸出特征,concat為拼接操作,FC為全連接層,conv為卷積層。使用多尺度特征融合可以進一步增強模型的表達能力,使其能夠更準確地進行分類。

3 實驗分析

3.1 數據集介紹

本文所用鋼鐵缺陷樣本圖像均來源于東北大學(NEU)表面缺陷數據庫中的NEU-CLS-64[9],該數據集收集了熱軋鋼帶的9種典型表面缺陷,即軋入氧化皮(RS)、斑塊(Pa)、裂紋(Cr)、麻面(PS)、夾雜物(In)、劃痕(Sc)、油污(Sp)、坑洼(Gg)和銹蝕(Rp)。NEU-CLS-64數據集中共有7 226張圖片且每個類別圖片數量不等,例如夾雜物(In)775張、坑洼(Gg)296張、油污(Sp)438張,但這些圖片的分辨率全部是64×64,這種低分辨率圖片無疑會給網絡正確分類帶來很大的難度。

3.2 實驗環境與設置

實驗使用基于NVIDIA 2080Ti GPU和Intel i7 9700K CPU的環境,采用PyTorch 1.8深度學習框架。訓練集與測試集之間的劃分比例為8∶2,即80%的數據用于訓練,而剩余20%的數據則用于驗證。在實驗中,所輸入的圖片分辨率為64×64,并采用大小為3×3的卷積核來構建Conv卷積塊,同時各個Conv卷積塊的通道數C按照64—160—320—640的順序遞增。為獲得具有更佳語義信息的特征圖,在鄰域自注意力模塊中全部采用點卷積技術。在優化器方面,使用SGD優化器,設置動量為0.9,學習率從0.01開始,衰減因子為0.05。在NEU-CLS-64數據集上進行了100次epoch訓練,每批次訓練的樣本數量為64。在第80個和90個epoch之后,采用學習率下降策略,將學習率減少0.1個因子。

3.3 算法對比分析

在實驗設置相同的情況下,對本文所提模型與其他神經網絡方法進行了比較,結果見表1。從實驗結果中可以看出,本文所提的鄰域自注意力網絡具有最好的分類性能。相比于ViT-B/16[10]、Swin_t[11]、ResNet50[12]、MobileNet_v3_small[13]、DenseNet121[14]和EfficientNet_b2[15]分別在準確率上高出9.39%、5.11%、4.83%、3.30%、3.24%和2.97%。此外,在運行時間方面,鄰域自注意力網絡也是最快的,這意味著本文模型能夠在相同時間內處理更多的圖像,并且保持最佳準確率。在實驗過程中發現,最新的ViT-B/16和Swin_t取得了最差的效果。這兩個模型不僅計算量大、參數較復雜,而且性能表現較差。這是因為輸入圖像尺寸較小時,劃分出來的圖像塊會比較小,每個圖像塊中包含的信息量有限,導致分類器無法捕捉到足夠的信息,從而影響模型的性能。

表1 與其他方法在NEU-CLS-64數據集上的比較結果

3.4 可視化分析

圖3呈現了各個模型的熱力圖可視化結果。從圖中可以看出,ViT-B/16和Swin_t權重的關注點比較混亂,而鄰域自注意力網絡可以過濾掉其中一些不相關的區域,并將注意力集中在更重要的圖像特征上。尤其是在Cr、In、Pa、Rp和RS這些缺陷上,相較于其他方法,本文方法能更精確地定位缺陷位置。同時,鄰域自注意力網絡融合了不同尺度的特征,使得它更有可能在復雜環境中學習到有用的特征。

圖3 不同模型可視化的結果

4 結束語

本文提出了一種基于鄰域自注意力的鋼鐵表面缺陷分類算法,通過鄰域自注意力模塊定位目標對象的重要區域并增強特征的可識別性。為了保持特征圖信息完整,使用多尺度特征融合的方法融合4種不同尺度的自注意力特征并用軋入氧化皮(RS)、斑塊(Pa)、裂紋(Cr)、麻面(PS)、夾雜物(In)、劃痕(Sc)、油污(Sp)、坑洼(Gg)和銹蝕(Rp)9類鋼鐵缺陷圖像進行了測試。實驗結果表明,本文所提模型在低分辨的鋼鐵缺陷圖像中具有良好分類性能,可進一步獲取缺陷的位置信息,對鋼鐵表面缺陷分類具有重要的實際應用意義。