基于樂高采樣的自監督表示學習方法

許亦博 ,趙文義 ,李靈巧 ,楊輝華

(1.桂林電子科技大學 計算機與信息安全學院,廣西 桂林 541004;2.北京郵電大學 人工智能學院,北京 100876)

由于無法獲取較多的帶標簽的樣本,自監督表示學習在自然語言處理[1]和計算機視覺[2]的諸多任務中發揮越來越重要的作用。自監督表示學習是使用無標簽的數據,利用對比學習,設計不同的代理任務或使用混合任務以優化特征提取器,并將特征提取器應用到下游任務中。由于自監督學習在特征表示學習階段不需要標簽信息,因此在自然場景圖像處理、醫學圖像處理、遙感圖像處理、工業檢測等領域有著非常廣泛的應用。

目前,自監督表示學習算法主要包括對比學習[3-4]、代理任務[5-6]和混合監督任務3類研究方法。對比學習算法的主要思想是先使用特征提取器提取一個樣本的不同視角的特征向量,然后建模2個向量的相似性與差異性,最后使用信息-噪聲對比估計[7](information-noise contrastive estimation,簡稱InfoNCE)損失函數優化特征提取器?；谠摾碚?He等[3]提出基于動量更新的對比自監督表示學習方法(MoCo),該方法使用記憶庫和孿生網絡結構,通過使用記憶庫的方法將特征向量存儲起來以增加對比樣本的數量,然后使用孿生網絡的結構對同一樣本的2個視角進行特征提取,最后通過建模2個向量的相似程度來優化模型。由于使用記憶庫存儲特征向量增加了存儲空間,Chen等[8]通過改進動量更新的對比學習,簡化復雜的結構并且去掉了記憶庫,提出了一種自監督表示學習的簡單結構(SimCLR)。而上述方法均無法獲取圖像的細節信息,因此,Wang等[9]提出一種基于密集映射用以保留局部細節信息的方法,利用密集預測的特性,將提取到的特征圖分成若干個特征向量,并分別對比每個特征向量。由于每個特征向量對應不同的細節信息,該方法能在密集預測的下游任務(目標檢測)中取得較好的效果,但對全局預測的下游任務(場景識別)并無幫助。上述基于對比學習方法依賴于大量的數據和較多的訓練周期,并且無法同時獲得局部信息和全局信息,導致數據不能充分利用。

基于代理任務的方法主要是通過人為設計一種監督學習的方法,根據設計的方法使用無標簽數據做監督學習,從而優化特征提取器。如基于補丁排序[10]的方法和預測旋轉角度[11]的方法被用來優化特征提取器。最近,He等[6]引入了基于非對稱編解碼網絡(MAE),通過預測被掩蓋的區域來優化模型,這種方法能獲得較為魯棒的特征提取器,但是需要使用計算量和參數量都較大的Transformer網絡[12]才會有效果,加重了計算和存儲負擔。然而這些基于代理任務的方法在使用CNN 作為特征提取器時,收斂速度較慢且無法提取細節信息,導致在下游任務中表現較差。

基于混合監督任務的方法目前很少有研究涉及,Chen等[13]提出一種聚類和定位共同優化的混合監督方法,這種方法雖然能加快收斂速度,但細節特征的提取能力不強。

針對現有方法存在的收斂速度慢和無法獲取細節特征的問題,提出了一種基于樂高采樣的自監督表示學習方法。本文采用了混合監督任務的方法。整體模型中除了引入樂高采樣模塊和特征提取模塊外,在混合優化模塊還引入聚類分支、定位分支和細節對比分支,通過對不同補丁的聚類和定位分支來恢復原始圖像,并學習全局特征;通過細節對比分支來學習細節特征,從而使模型在更多的下游任務中具有較強的適應性。

1 基于樂高采樣策略的自監督學習模型

基于樂高采樣策略的自監督學習模型如圖1所示。對輸入的任意大小的無標簽圖像,本方法首先引入樂高采樣策略,生成一系列如圖1所示的5個補丁組,其中4個補丁相互鄰接,補丁5與這4個補丁均有重疊。將殘差網絡作為特征提取器對樂高采樣補丁進行特征提取,輸出對應的特征向量,其中每個特征向量代表一次樂高采樣的補丁;然后對這些特征向量使用不同的損失函數進行優化。

1.1 樂高采樣

由于在自監督表示學習的過程中需要大量的正負樣本對以維持較高的性能,傳統的采樣方法在計算資源有限的情況下無法構造出較多的正負樣本,考慮到特征提取器的作用只是將原始樣本轉換為低維的特征向量,因此可以使用通過拼接后的樣本進行特征提取來增加樣本數量,并以此來學習樣本間的關系。假設有n個從無標簽數據集中隨機選擇的單張圖片樣本X={x1,x2,…,xn},對一張輸入圖片,樂高采樣模塊輸出m個大小為112×112的補丁,經過特征提取模塊的提取,生成m個對應的特征向量,這些特征向量經過不同的映射后用來優化模型。在該方法中,每個樣本都會經過樂高采樣器,所以共生成nm個分辨率為112×112的補丁。隨后將這些補丁隨機打亂,并將每4個補丁拼接成一個224×224的完整樣本,重新組合形成新的數據集,再將這些構造的樣本送入特征提取器。樂高采樣的結構如圖2所示,圖中藍色標號5的補丁不僅能增加樣本數量,還能促使模型學習更多的細節信息。

圖2 樂高采樣示意圖

1.2 特征提取

為了公平對比,該方法的特征提取模塊和其他方法的特征提取模塊保持一致,使用ResNet-50 網絡[14]作為特征提取器。該網絡包含4個階段,每階段中卷積塊的數量分別是3,4,6,3。每個卷積塊包含一個1×1通道下降的卷積,一個通道不變的3×3卷積和一個1×1通道數上升的卷積,2個1×1卷積的作用是為了保持輸入通道和輸出通道數的一致,并降低3×3卷積處的計算量。除此之外,為了適應自監督學習,把網絡最后的全局平均池化層和全連接層去掉,并使用混合優化模塊代替。

1.3 混合優化

在混合優化模塊中,使用全局聚類分支,全局定位分支和局部細節對比分支共同優化。用到的損失函數分別為有監督聚類損失函數、交叉熵損失函數和對比損失函數。

1.3.1 全局聚類分支

全局聚類分支的作用是將來自同一張原始圖片的樂高采樣并打亂后的數據聚類,該分支采用信息-噪聲對比估計(InfoNCE),定義為將每個樣本的采樣結果聚類到一起,它最大化正負樣本之間的距離,同時最小化正正樣本之間的距離。在本實驗中,將來自于同一個樣本的采樣結果視為正樣本,將來自不同樣本的采樣結果視為負樣本,根據以上分析,全局聚類分支的損失函數表示如下:

其中ls表示針對某一個特征向量的損失函數,

式中ls,j表示信息-噪聲對比估計損失函數,可以用式(3)表示:

1.3.2 全局定位分支

全局定位分支的作用是輔助全局聚類分支將樂高采樣并打亂后的數據恢復成原始狀態。采用交叉熵損失函數(cross entropy,簡稱CE)進行優化,目的是用來判斷每個補丁來自原圖的哪個部位,輔助恢復原始圖像。其損失函數如下:

式中:C表示經過特征提取后的特征向量集合;Cgt表示根據特征向量的相應位置生成的偽標簽。

1.3.3 局部細節對比分支

細節對比模塊優化的目標是將樂高采樣的前4個補丁生成特征向量的加權平均與第5個補丁生成的特征向量進行對比,這樣操作不僅能促使模型學習4個模塊中語義一致性的特征,也能使模型更關注于細節特征,細節對比模塊的操作如圖3所示。使用傳統的對比損失函數(CON)進行優化,損失函數表示為

圖3 局部細節對比分支

其中:Vi表示前4個向量的加權平均;Vh表示第5個向量。

最終整個模型的損失函數為3個損失函數的加權平均,定義如下:

模型優化的目標是最小化L。

2 實驗結果及分析

為了驗證提出方法的有效性,首先建立適用于自監督學習的數據集并對模型進行自監督表示學習,然后固定自監督學習階段學習到的模型參數,將模型遷移到線性分類和目標檢測的下游任務中進行測試。其中在自監督學習階段使用的是無標簽數據,在下游任務中使用的是有標簽數據。本文使用提出的算法在線性分類和目標檢測等下游任務中與4種經典的自監督算法進行對比,這些算法包括基于動量更新的對比自監督表示學習算法(MoCo)、簡化的自監督學習算法(SimCLR)、改進的基于動量更新的對比自監督表示學習算法(MoCo-v2)[15]、基于拼接圖像聚類的自監督表示學習算法(JigClu)。為了公平對比,使用相同的訓練數據和參數對各個模型進行訓練。仿真實驗平臺采用AMD EPYC 7302 16核CPU,內存為128 GiB,顯卡型號為GeForce RTX3090 24 GiB,深度學習框架為PyTorch。

2.1 數據集構建與模型的訓練

為了驗證提出模型的有效性,使用ImageNet[16]、CIFAR[17]和若干細粒度分類數據集作為實驗數據。其中,ImageNet數據集源于自監督訓練階段;CIFAR數據集用來快速驗證提出模型的遷移性能;細粒度分類數據集用來驗證模型對細節特征的提取能力。

ImageNet數據集常被用來驗證模型的特征提取能力,是一個非常有挑戰性的數據集,它包含128萬張訓練樣本、50 000張驗證樣本,具有1 000個類別。CIFAR數據集作為計算機視覺最經典的數據集之一,常用來快速驗證模型的性能和做簡單的消融實驗,CIFAR數據集包含50 000張訓練樣本和5 000張驗證樣本,其中CIFAR-10和CIFAR-100兩個版本分別包含10和100個類別。

用3種常見的細粒度分類數據集Caltech-UCSD Birds (CUB-200-2011)[18]、Stanford Cars[19]和FGVC Aircraft[20]來驗證模型對細節特征的提取能力。這些數據集的統計信息如表1所示。

表1 數據集統計信息

在自監督訓練階段,設置批大小為256,初始化學習率為0.03,并且使用隨機梯度下降的方法(stochastic gradient descent,簡稱SGD)優化模型參數,其中動量和權重衰減值分別設置為0.9和0.000 1,自監督階段模型一共被訓練了200個周期。在線性評估階段,批大小被設置為256,初始學習率設置為10.0,并在第60和80周期衰減到原來的0.1倍,訓練的總周期數為100。

2.2 客觀評價結果

表2為本方法和幾種經典方法在ImageNet的線性分類任務中對比的結果。

表2 分類任務結果

表2表明本方法在下游的線性分類任務中具有較高的準確率。此外,相比于MoCo、Moco-v2、Sim-CLR等基于孿生網絡的結構,本方法在保持較高性能的同時節約了大量的計算量。

為了證明提出的模型有較強的魯棒性和泛化性,同時具有較強的細節提取能力,在CIFAR數據集上和3種細粒度分類數據集上進行實驗,結果如表3、4所示。

表3 CIFAR數據集實驗結果

表3中:finetune代表使用較小的學習率更新權重;linear表示凍結權重,只更新線性層。而random init表示權重隨機初始化,supervised表示有監督訓練,Lego表示使用提出方法的預訓練模型初始化權重。

從表3、4可見,本方法在CIFAR數據集表現出較強的遷移學習能力,在3種細粒度分類數據集中表現出對細節特征有更強的提取能力。

綜上可得,本方法得到的預訓練模型相較于MoCo、MoCo-v2、JigClu等幾種經典方法具有更強的魯棒性和泛化性。

2.3 主觀評價結果

為了能直觀地表示提出方法的有效性,使用類激活圖(class activation mapping,簡稱CAM)來直觀地展示提出模型的魯棒性和對細節特征的提取能力。

圖4表示使用預訓練權重在不同數據集上生成的CAM。在CAM 中,紅色區域表示模型重點關注的區域,也即激活區域,藍色區域表示未激活區域。觀察圖4可以直觀地看出本方法在應對不同形狀、不同尺寸的全局特征,難以辨識的細節特征時能更好地激活相對應的區域。在圖4中,(a)是原始圖像,(b)和(c)是MoCo在layer3和layer4的CAM,(d)和(e)是MoCo-v2在layer3和layer4的CAM,(f)和(g)是JigClu在layer3和layer4的CAM,(h)和(i)是本方法在layer3和layer4的CAM。

3 結束語

提出了一種基于樂高采樣的自監督表示學習方法。首先,利用樂高采樣器對無標簽數據進行采樣并形成新的數據集;然后使用特征提取網絡對新的數據集進行特征提取,并生成對應的特征向量;最后,使用不同的損失函數對不同的分支進行優化,以得到魯棒性和泛化性較強的特征提取器。理論分析和實驗結果表明,所提方法在客觀評價指標和主觀評價指標上皆獲得較好的效果,能使自監督學習的下游任務得到更大的提升。

所提方法在自監督訓練方法只是單純使用RGB圖像作為輸入,目前很多實際應用場景中,除了RGB圖像,各種模態的信息如文本、音頻等也有著重要的使用價值,多模態學習是研究的熱點之一。在未來,如何將自監督學習與多模態學習結合到一起,是一個值得深入的研究方向。