基于對比學習的小樣本金屬表面損傷分類

吳冠榮，李元祥，王藝霖，陸雨寒，陳秀華

（上海交通大學航空航天學院，上海 200240）

0 引言

對金屬表面存在的損傷進行分類是金屬類工業產品質量檢驗和反饋工作中不可或缺的一環，通常需要有專業知識的人借助電鏡等器材進行判斷［1］。但是損傷類型多、視覺疲勞、工業流水線的壓力等因素容易影響主觀判斷，人工分類的準確率忽高忽低。隨著機器視覺技術的發展，神經網絡對視覺特征的強大學習能力可以彌補人工分類的不足。

基于機器學習的金屬損傷分類按照是否使用初始化參數，分為遷移學習方法和深度學習方法。遷移學習方法希望從ImageNet 等大型訓練集中學習到最通用的特征，在目標數據集上進行微調［2］。而深度學習方法通過調整模型參數來擬合當前數據的特征分布［3］，或者通過加深模型層數，提高模型對多維特征的表達能力［4］。但金屬損傷的紋理特征和常見的大型數據集差異很大，公開的預訓練模型遷移效果有限，并且金屬損傷種類繁多，一些偶發性、新產生的損傷不滿足神經網絡對訓練數據量的需求，若直接使用神經網絡進行有監督訓練，模型易于過擬合、泛化性能差。因此，在很多應用場景下需要考慮小樣本問題。

較早的小樣本問題解決方案通常是學習合成偽數據［5］。這類方法依賴于擴充樣本的質量，如果合成器效果一般，則容易引入噪聲數據，效果有時不升反降［6］。目前，小樣本分類研究聚焦于如何學習1 個更好的特征表示，使下游任務中即使只有少量標注信息，分類模型也能表現良好［7］。常見方法包括元學習框架［8-10］、度量學習［11-12］、知識蒸餾［13-14］等，目的都是為了找到更好構造特征空間的方式。ProtoNet［8］使用度量手段進行原型學習，由于其具有一定的靈活性和高效性，因此仍廣泛應用在目前算法中。文獻［9］在通用的兩階段微調訓練機制基礎上，增加了一步預訓練，利用DINO 機制在額外的無標簽數據上訓練主干網絡的特征提取能力，再在ProtoNet 上訓練模型，實現了小樣本分類任務的SOTA（State-Of-The-Art）效果。

文獻［15］首先將小樣本解決方案引入金屬損傷分類，利用生成對抗網絡（GAN）生成大量無標注的鋼材表面損傷樣本，在擴充樣本量后，再用殘差網絡進行分類。文獻［16］利用Wasserstein 距離的生成對抗網絡（WGAN）有效擴充鋼板氣泡損傷，提高檢測精度。然而，生成偽樣本的方法容易模式固化，尤其對于金屬表面損傷問題，金屬材料種類多，損傷種類更多，這些方法僅擴充某種材料或某類損傷的數據量，因此算法的適用性不高。文獻［17］針對鋁材常見損傷類別各提取一種特征圖CCMs，并利用注意力模塊生成小樣本類別的建議特征圖（RCMs），對CCMs 和RCMs 進行映射和轉化，分類結果由CCMs和RCMs 共同決定。該方法基于原型學習思路，將類內共性作為區分標準，預測時衡量當前樣本與各個類別原型特征的相似度，最相似的類別作為分類結果，構造原型的范式在一定程度上緩解了模型泛化性差的問題。但是，金屬損傷紋理細微、多變決定了特征分布的類間距離小、類內距離大，現有方法局限于在弱區分度的特征分布基礎上根據類內共性完成分類，導致少樣本、新類別分類精度低。

為此，本文提出一種基于對比學習的小樣本金屬表面損傷分類方法CLFS。該方法在元學習機制下，以外層模型對比、內層模型度量作為元任務，引入類別標簽嵌入作為外模型對比不同類別圖像特征的監督信號，使得在新的特征空間下，類內特征距離近，類間特征距離遠，視覺特征相似的類別中心距離近。在金屬損傷分類任務中，外模型對比學習代理任務可以間接使得原有特征空間中不同損傷類別之間的差異更加明顯，從而增強了內模型的度量學習能力，有效提高分類精度；引入類別信息作為監督信號可以減少對比學習過程中樣本噪聲的影響，提高小樣本原型向量估計的魯棒性，從而提升對新類別的泛化能力。

1 小樣本學習描述

關于小樣本問題的任務設定，本文采用最常見的目標因子分析（TFA）數據劃分方法（TFA-split）和全類別測試方案［18］。TFA-split 將數據集分為訓練集和測試集：訓練集選取具有較高資源的基礎類別（Cbase），共Mb類；測試集包含訓練集中未出現的新類別（Cnovel），新類別是低資源的小樣本問題，新類別數共Mn類。TFA-split 在迭代過程中保持訓練集和測試集固定不變。不同于僅測試Cnovel的狹義小樣本問題，本文同時測試模型對Cbase和Cnovel的分類性能。

本文模型的訓練采用元學習模式。元學習模式在訓練過程中將傳統的單個循環多批次的訓練任務拆分成數個元任務。每個元任務目標一致、構成一致，從所有類別中選取K個類別，每個類別各選N個樣本作為元任務輸入。1個元任務又稱K-wayN-shot任務，1 個元任務的工作周期為1 個episode。在訓練過程中多次迭代元任務使得模型滿足最終的小樣本測試任務要求。元學習模式實現了從以數據為中心向以任務為中心的轉變，從而減少對樣本量的依賴，并提高小樣本學習的性能。

本文根據資源高低將1 個金屬損傷樣本集按類別區分為Cbase和Cnovel，并將Cbase樣本劃分出訓練集Dtrain和測試子集樣本作為測試子集則測試集樣本為由于所有樣本需要被統一為元任務的輸入格式，因此按照K-wayN-shot格式重構Dtrain和Dtest。

在訓練階段，每個元任務先從Cbase中抽取K（K≤Mb）個類別，針對抽出來的每個類別對應從Dtrain中抽取N張圖片作為訓練樣本，稱這K×N張樣本為1個元任務的支持集S。對同樣的K個類別再從Dtrain取N張不同的圖片作為驗證樣本，稱為1 個元任務的查詢集Q。S和Q的類別相同、樣本不同，合在一起作為1 個元任務的輸入樣本。元學習過程是指在每次元任務中利用S的信息對Q進行分類，從而更新模型參數。Dtrain被重構成m個元任務，即Dtrain=本文采用2 種常用的K、N設定，即5-way 5-shot 和5-way 1-shot。

在測試階段，將Dtest按照K-wayN-sho（tK≤Mb+Mn）格式重構成多個元任務，模型在多個元任務上的平均分類性能即最后的性能。

2 網絡架構

雖然元學習模式能夠讓模型快速適應新的任務，但是因金屬損傷的類內方差大、類間相似度大等特性，使得很多小樣本學習方法的效果明顯下降。為提高模型的泛化能力，本文提出引入類別標簽信息作為對比學習的監督信號。利用對比損失函數懲罰不同類別特征之間的距離，能有效地劃分外模型特征空間的類間區域，間接優化內模型的特征分布，使得內模型在分類時可以不受無規律特征分布的影響。

2.1 訓練模型

CLFS 分類模型在訓練階段分為內模型和外模型，內模型負責圖像分類，外模型負責優化特征分布的代理任務，內外模型聯合訓練。圖1 所示為在K-wayN-shot設定下CLFS 的網絡架構，原始樣本經過預處理分為圖像輸入和類別標簽嵌入作為整個訓練模型的輸入信號。圖1中Query set即所須訓練的樣本編碼得到的向量集合。內模型主要包括利用ResNet 18的特征提取模塊F和原型向量（包括支持向量和查詢向量）生成模塊。外模型除了與內模型共用的特征提取模塊F之外，主要由類別標簽嵌入、特征轉譯（特征編碼和類別標簽編碼）和對比分支3個部分組成。

圖1 在訓練階段K-way N-shot 設定下CLFS 的網絡結構Fig.1 Structure of CLFS network with the K-way N-shot setting in the training phase

2.1.1 內模型

本文使用ResNet 18作為特征提取模塊，ResNet 18由4層殘差模塊構成，選擇特征層8倍下采樣結果作為F的輸出向量。支持集樣本和訓練集樣本分別經過F映射得到支持向量和查詢向量，以計算原型向量集，并采取度量學習的方式完成分類任務［8］。

元任務的內模型工作流程的輸入為X：K×N張支持集圖片；X'：K×N張查詢集圖片，輸出為K*查詢集分類結果。工作流程如下：

1）第k類的N張支持集圖片Xk經過特征提取模塊F的投影后，形成支持向量子集Sk=F(Xk)；

3）查詢集圖片X'經過特征提取模塊F的投影后，形成查詢向量集

4）根據S計算每個類別的原型向量

5）計算查詢向量和每個原型向量之間的歐氏距離Dk=‖Q-ck‖2；

6）第i個查詢向量qi距離最近的原型向量所代表的類別為qi對應查詢圖片分類結果

2.1.2 外模型

在訓練階段增加的外模型中，利用類別標簽嵌入信息作為確定的類別中心，減少對比學習過程中樣本噪聲的影響，從而間接提高內模型原型向量估計的魯棒性。

1）類別標簽嵌入。CLFS 所用的3 個數據集標簽如表1 所示。本文的類別標簽嵌入流程：當前數據集的所有類別標簽經過獨熱編碼后，按照類別順序存儲為Hash 表，在每次元任務中對用到的類別進行Hash 映射，得到初始類別嵌入向量。

表1 數據集標簽信息Table 1 Label information of datasets

2）特征轉譯。該部分包括2 個模塊：（1）特征編碼器，通過特征提取模塊得到的結果向量在進入對比空間前也須經過1 層多層感知機（MLP）映射；（2）類別標簽編碼器，在預處理階段嵌入的初始類別標簽向量先經過1 層可學習的MLP 映射，再進行有監督對比學習。本文設計2 個特征轉譯模塊主要是為了解決在對比學習中不同模態之間的差異性問題，提升對比特征空間的一致性，同時，構建輔助特征空間使得對比學習過程不影響內模型的度量空間。此外，在特征轉譯模塊提供的新特征空間下，對比學習的動態性［22］讓類別標簽嵌入和圖像特征互相影響。類別標簽嵌入的結果作為類別中心監督特征分布，通過反傳損失給特征編碼器、特征提取模塊逐步優化分類邊界。圖像特征信息通過反傳給可學習的類別標簽編碼器，類別中心的分布位置可以繼承視覺特征相似性。類別標簽監督不同類別之間的特征分布情況如圖2 所示。從圖像紋理層面鋼材油斑類別（oil_spot）和鋼材夾雜物（inclusion）損傷類別相較于油斑和折痕（crease）更相似，則訓練后油斑類別的特征分布更接近夾雜物而不是折痕。

圖2 類別標簽監督不同類別之間的特征分布情況Fig.2 Features distribution among different categories in category labels supervise

3）對比分支。對比學習過程結合2 個可學習的特征轉譯模塊，在訓練過程中拉近2 個新映射的特征空間距離，可以有效減小不同模態特征之間的差異。經過相同結構的特征轉譯模塊，類別標簽和圖像2 種模態的特征被映射到同一空間，主要目標是在新的特征空間下，利用離散的類別標簽先驗信息對圖像特征分布進行有監督對比學習，使得同類別的特征分布聚合性更好，不同類別的特征區分度更高。對比分支示意圖如圖3所示。

圖3 對比分支示意圖Fig.3 Schematic diagram of contrasting branches

在特征轉譯后的圖像特征和類別標簽嵌入經過歸一化（Normalization）和點積（dot product）操作，得到余弦相似度矩陣。

2.1.3 損失函數

損失函數分為度量損失和對比損失，分別是內模型對小樣本元任務分類結果的監督和外模型對特征分布情況的監督。

度量損失Ldistance是指在1 次元任務的訓練過程中，攜帶真值的K×N張查詢圖片，經過特征提取形成查詢向量集對于查詢向量qi，計算其與對應真值類別原型之間的Softmax 負概率分布，即懲罰當前所得特征與真值類別的原型向量之間的距離［8］：

受SimCLR［22］的啟發，本文選取余弦相似度作為轉譯后的圖片特征和類別標簽特征之間的對比損失Lcontrast。圖3 中最終構成的余弦相似度矩陣分別從標簽特征維度和圖像特征維度懲罰同類別之間的距離，從而間接增大類間距離。假設1個episode由2N張樣本組成，包括N張支持樣本和N張查詢樣本，標準化后標簽特征向量表示為vj，對應真值為。查詢圖片和支持圖片所得到的標準化向量均表示為wi，對應真值為，l 表示當｛｝中條件為真時，則取值為1，反之為0。則1 個episode 的對比損失表達式如下：

其中：溫度系數τ用來調節對比學習中的分類細粒度，τ越大，對不同樣本間差異性的重視程度越高。

總的損失函數表達式如下：

通過一定的權重系數α將內外模型的損失函數進行統一，外模型的轉譯層參數由外模型對比損失單獨更新，其他模塊的可學習參數由整體損失函數共同優化。

2.2 預測模型

經過訓練后，多模態對比學習模塊加強了內模型對不同類別特征分布的區分能力，元學習模式也使得內模型提前適應小樣本任務。因此測試階段只需要內模型進行前向推理，不需要外模型進行監督和輔助。整個測試推理過程與訓練階段內模型的1 次前向傳播步驟一致，1次K-wayN-shot 的全類別測試任務包括訓練階段未涉及的新類別Cnovel，體現了CLFS 小樣本分類的泛化性能。預測模型架構如圖4 所示。預測模型為訓練模型的內模型，最終的分類結果k*由距離查詢向量最近的原型向量的所屬類別決定。

圖4 預測模型架構Fig.4 Framework of prediction model

3 實驗

3.1 數據集與實驗設置

本文實驗選取金屬表面損傷常用的3 個數據集：GC10 鋼板表面損傷數據集［19］、東北大學鋼材損傷分類數據集（NEU）［20］、天池鋁型材料表面損傷數據集（APSD）［21］。

GC10 數據集包括10 種類型的鋼材表面缺陷，每個類別的樣本數量不等，共3 570 張灰度圖。由于從生產線統一采集，因此GC10 樣本尺寸均為2 048×1 000 像素。NEU 數據集由東北大學機器視覺與機器人團隊發布，包括6 類常見鋼材表面損傷，每類有300 張灰度 圖，尺寸固定為200×200 像素。APSD 數據集來自天池2018 鋁材表面瑕疵識別比賽，本文選取其中9 個類別用于訓練，每個類別的樣本量從50 到200 不等，尺寸統一為256×256 像素。本文實驗將原始RGB 圖像統一轉為灰度圖。

為統一輸入圖像尺寸對骨干網絡的影響，在預處理階段將所有樣本保持初始的長寬比壓縮至150×150 像素左右，再用背景像素填充成方陣。CLFS 采用ResNet 18 作為圖像特征提取層，1 層MLP 作為轉譯層，所有網絡使用Adam 優化器，設定的初始學習率為0.001，每20 輪減小1 次學習率，設置對比損失中溫度系數τ=2，度量損失占整體損失的權重α=0.2。

實驗軟件配置包括Python 3.8.10、PyTorch 1.7.0、OpenCV 4.6、Numpy 1.21，硬件平臺為GPUNVIDIA GeForce RTX 3090 和CPU-Intel?Xeon?Platinum 8255C。

3.2 實驗結果與分析

本文重復10 組對比實驗，取所有分類精度結果的標準差作為偏差范圍。所選取的比較對象不僅包括傳統 的兩階 段微調（TFA）［23］，而且包 括MatchingNet［24］、RelationNet［25］、ProtoNet［8］等主流 的小樣本分類方法。同時，在全類別測試的基礎上，通過混淆矩陣比較新類的分類性能，并進行特征分布可視化分析。

3.2.1 分類精度對比

CLFS 在5-way 5-shot 和5-way 1-shot 設定下，全類別測試的分類精度都得到了較大幅度的提升，如表2 所示，加粗表示最優數據。

表2 不同方法的分類精度比較（全類別測試）Table 2 Classification accuracy comparison among different methods（full category test）%

在5-way 5-shot 設定下，CLFS 的分類精度明顯優 于TFA，而 且CLFS 比ProtoNet、MatchingNet、RelationNet 3 個主流小樣本方法的分類精度至少提高了5.24、1.39 和6.37 個百分點（與次優結果比），分類錯誤下降率分別為36.00%、17.94%和66.15%。對于5-way 1-shot 設定，CLFS 相較于3 個主流小樣本方法，分類精度提升效果更顯著，分別至少提高8.34、3.01 和4.61 個百分點（與次優結果比），分類錯誤下降率分別為28.32%、23.37%和46.57%以上。

3.2.2 混淆矩陣分析

在GC10 數據集5-way 5-shot 下ProtoNet 與CLFS 分類混淆矩陣如圖5 所示。前7 行（列）類別為此次實驗中設置的Cbase，最后3 行（列）為小樣本的Cnovel，混淆矩陣以行為單位包括真正例（TP）和假反例（FN）類結果，以列為單位包括TP 和假正例（FP）類結果。從全類別測試結果中可以看出，對于GC10 數據集，ProtoNet 不僅對Cnovel的分類效果較差，而且對Cbase漏報率和誤報率也較高。CLFS 在同樣情況下Cbase的FN 類、FP 類結果基本清零，即真值為Cbase的類別基本分類正確。特別地，CLFS 有效提高模型對小樣本Cnovel的泛化性，相較于ProtoNet，Cnovel分類精度由36.53%大幅提升至69.12%。針對NEU 數據集和APSD 數據集，Cnovel分類精度分別由82.43%提升至91.57%，31.89%提升至48.23%。

圖5 在GC10 數據集上5-way 5-shot 下ProtoNet 與CLFS 分類混淆矩陣Fig.5 Classification confusion matrix for ProtoNet and CLFS in 5-way 5-shot on GC10 dataset

3.2.3 特征分布可視化

為動態監督特征分布，本文采取t-SNE 對統一Backbone 后的ProtoNet 和CLFS 的同一特征層進行降維映射。在GC10 數據集5-way 5-shot 下ProtoNet和CLFS 的特征分布對比如圖6 所示（彩色效果見《計算機工程》官網HTML 版）。從圖6（a）～圖6（c）可以看出，ProtoNet 的特征分布中大多數類別的特征點互相穿插，Cbase和Cnovel都沒有區分度。圖6（d）～圖6（f）所示為同樣條件下CLFS 在測試過程中前向預測得到的特征分布，類別之間的間距明顯、類內聚合度也顯著提高。特別地，將Cbase和Cnovel分開分析，Cbase的特征邊界清晰，類別分布的相鄰關系保留了視覺特征的相似性。由于測試階段才引入Cnovel，因此CLFS 模型對于Cbase和Cnovel之間的區分度下降，但是與圖6（c）的Cnovel分布對比可以看出，CLFS 有效提高了類內聚合性。實驗結果表明，CLFS 利用外模型雙模態對比學習，有效地解決了類間相似性較高的金屬損傷分類問題。

圖6 在GC10 數據集5-way 5-shot 下ProtoNet 和CLFS 的特征分布對比Fig.6 Comparison of feature distribution between ProtoNet and CLFS under 5-way 5-shot for GC10 dataset

3.3 消融實驗

為了分別驗證對比學習和引入類別標簽信息能否提升模型對小樣本的分類能力，本文使用ProtoNet 作 為baseline，拆 分CLFS 的外模型為baseline+對比損失、baseline+標簽嵌入，在3 個數據集上完成5-way 5-shot 的消融實驗。

消融實驗結果如表3 所示。加入對比損失和引入類別標簽嵌入的方法相較于baseline 分別提升了1%～2%的分類精度，驗證了對比損失項和類別標簽嵌入的有效性。CLFS 所設計的外模型通過構造輔助特征空間，將對比損失和標簽嵌入2 種方法組合，對分類效果有更顯著的提升，單個數據集的分類精度相較于baseline 分別提升5.24、3.03 和14.00 個百分點，分類錯誤下降率分別為36.00%、32.27% 和81.11%。CLFS 的內外模型機制具有一定的有效性。

表3 消融實驗結果Table 3 Results of ablation experiment %

3.4 超參數分析

CLFS 涉 及2 個超參數τ和α，GC10 數據集 在5-way 5-shot 設定下進行超參數分析，根據經驗［22］確定超參數τ和α的取值范圍分別為［0.1，1，2，5］和［0.1，0.2，0.5，2］。本文使用網格法對每組參數重復3 次實驗。超參數對模型性能的影響如表4 所示，在合適的超參數范圍內，分類性能對超參數的敏感度較低，超參數優化空間較為平緩，峰值處的誤差范圍更小，因此魯棒性更好。從表4 可以看出，當τ=2、α=0.2 時，模型性能最佳。

表4 超參數對模型性能的影響Table 4 The impact of hyperparameters on model performance %

4 結束語

針對小樣本金屬損傷特征分布聚合性低、新損傷類別泛化性差的問題，本文通過設計內外模型機制，提出一種基于對比學習的小樣本金屬表面損傷分類方法。內模型以傳統度量學習為目標，完成小樣本分類的任務；外模型構造輔助特征空間，在該空間下利用類別標簽嵌入信息，動態對比不同類別的特征分布，從而實現聚合同類別特征，得到明確清晰的分類邊界，提高模型對base 類別的分類精度，而且在較大程度上提升對novel 類別的泛化能力。在3 個常見的金屬損傷數據集上驗證本文方法的有效性，特征分布的可視化進一步展現了本文方法在增加類間距離、聚合類內特征的任務上具有顯著的提升效果。目前，小樣本金屬損傷分類研究尚處于探索階段，可參考的數據集和實驗結果有限。本文僅初步探索了對比學習結合多模態優化特征分布在小樣本分類算法方面的可行性，下一步將完善標簽嵌入模型，便于從語義相似性角度補充小樣本金屬損傷的先驗信息，進一步優化分類性能，同時使得結果具有可解釋性。