結(jié)合注意力轉(zhuǎn)移與特征融合算法的在線知識蒸餾

梁興柱,徐 慧,胡 干

(1.安徽理工大學 計算機科學與工程學院，安徽 淮南 232001；2.安徽理工大學環(huán)境友好材料與職業(yè)健康研究院(蕪湖)，安徽 蕪湖 241003)

近年來，深度神經(jīng)網(wǎng)絡憑借強大的特征學習能力在計算機視覺處理中取得了令人欣喜的成績[1-2]。但是，功能強大的模型往往會伴隨大量的參數(shù)，占據(jù)較大的內(nèi)存，不利于模型的部署與應用[2]，因而提出了深度學習模型壓縮技術，包括剪枝、二值化、輕量化模型設計[3-4]、知識蒸餾(Knowledge Distillation , KD)[5]等。其中，知識蒸餾又可以劃分為離線知識蒸餾(Offline KD)和在線知識蒸餾(Online KD)。傳統(tǒng)Offline KD[5]是一個兩階段蒸餾方法，必須先預訓練一個功能強、參數(shù)多的教師模型，然后再將教師模型學到的知識遷移到性能較弱、參數(shù)較少的學生模型上，從而得到一個速度快、能力強的網(wǎng)絡，達到減少參數(shù)、提高學生模型性能的目的，但存在訓練時間長、計算成本高、占據(jù)內(nèi)存大等缺點。Zhang等[6]利用深度相互學習(Deep Mutual Learning,DML)使學生模型直接從其他學生模型的預測中學習。Lan等[7]提出采用一個門控單元作為網(wǎng)絡共享低層，在動態(tài)建立教師模型的同時訓練多分支網(wǎng)絡。Romberg等[8]進一步優(yōu)化了FitNet 網(wǎng)絡并率先在知識蒸餾領域提出利用教師與學生間的注意力圖學習來代替特征圖學習。Chen等[9]在OKDDip兩級蒸餾方法的研究基礎上使用分類器多樣化損失函數(shù)和特征融合模塊來提高學生模型的多樣性和網(wǎng)絡中注意機制的性能。

現(xiàn)有的大部分方法很難構建一個功能強大的教師角色，且忽略了單個子分支性能的自我提升。因此，本文提出一種結(jié)合注意力機制(Attention Mechanism，AM)與特征融合(Feature Fusion Module，F(xiàn)FM)的在線知識蒸餾方法(KD-ATFF)，利用特征融合構建強大的教師角色指導模型訓練，同時將深層神經(jīng)元的注意力轉(zhuǎn)移到淺層網(wǎng)絡，進一步提升子分支的性能。

1 KD-ATFF

KD-ATFF擁有n個分支在線集成網(wǎng)絡模型，每個分支網(wǎng)絡由M個block組成。在沒有教師模型指導的前提下，每個block在訓練過程中將自己的特征圖轉(zhuǎn)化為注意力圖，模塊之間相互學習彼此的注意力地圖，增加知識差異性。各個分支由淺層網(wǎng)絡到深層網(wǎng)絡相互學習得到的多樣性知識被保留到最后1個block的特征，然后再將各分支的最后1個block的特征送入特征融合模塊進行信息融合，最后根據(jù)模塊內(nèi)部為各分支分配的權重組成集成教師指導各子模型訓練。KD-ATFF模型結(jié)構如圖1所示。

圖1 KD-ATFF模型結(jié)構

1.1 CL模塊

(1)

(2)

采用L2范數(shù)作為2個模塊進行互學習的損失函數(shù)，將2個模塊的注意力圖進行二次互學習得到的值進行相加后取1/2作為一次損失。模塊間的互學習目標函數(shù)為：

(3)

1.2 特征融合模塊

KD-ATFF共有n個相同的網(wǎng)絡架構分支，多分支的特性比來自單分支的特性包含的信息要豐富得多。為便于表示，所有的學生模型從1到n進行索引。由于網(wǎng)絡的深層會產(chǎn)生更豐富的語義信息，故將來自多個分支的最后1個block的特征作為特征融合模塊的輸入。這樣可以利用高級的語義信息來豐富特征，所產(chǎn)生的權值能取得較好的效果。加權集成目標ze可表示為：

(4)

式(4)中，f(·)為特征融合模塊中心塊的函數(shù)，為每個分支輸出相應的重要性分數(shù)；Fa為來自第a個分支的最后一個塊的特征圖；za為來自第a個分支的logits。

以3個分支作為輸入為例，來自每個分支的最后1個block的特征映射Fa將被連接在一起，然后送入中心卷積塊。中心卷積塊是由多個卷積層、批處理歸一化和ReLU激活函數(shù)組成，中心塊的最后一層是全連接層，用于融合來自多個分支的語義信息。與其他方法相比，特征融合模塊可以獲得更多的語義信息，能夠有效地提高模塊的性能。最終目標由各輔助分支的logits輸出zi加權和得到。特征融合模塊結(jié)構如圖2所示。

圖2 特征融合模塊結(jié)構

1.3 蒸餾模型的損失函數(shù)

1.3.1傳統(tǒng)標簽損失

(5)

式(5)中，T為溫度參數(shù)，取T=3。利用最小化交叉熵訓練，得到標簽學習的損失為：

Llabel=-∑ieilogqi

(6)

式(6)中，ei是標注的標簽分布；qi是最小化預測的類概率。

1.3.2注意力轉(zhuǎn)移損失函數(shù)

(7)

式(7)中，αm和αm+1分別為第m以及第m+1個block輸出的注意力圖。

1.3.3集成教師損失函數(shù)

KD-ATFF每個分支不僅從地面真實標簽中學習，還從通過特征融合模塊獲得的加權集成目標中學習。知識轉(zhuǎn)移是通過將學生模型生成的概率分布q與目標分布z對齊實現(xiàn)的。用KL散度表示其損失函數(shù)，第a個學生模型的預測分布為qa(a=1,2,…,n),每個輔助分支學習ze中提取的知識，故所有分支的蒸餾損失為：

(8)

則，整個KD-ATFF的損失函數(shù)為：

(9)

Llabel是第α個子模型的傳統(tǒng)知識蒸餾損失；θ，β是調(diào)節(jié)軟硬標簽比例的超參數(shù)。

1.4 KD-ATFF的算法流程

與兩階段傳統(tǒng)蒸餾訓練不同，Online KD中學生網(wǎng)絡和集成教師同時進行訓練。在每個子網(wǎng)絡進行相同的隨機梯度下降，并訓練整個網(wǎng)絡直到收斂，作為標準的單模型增量批處理訓練。批處理貫穿整個訓練過程，在每個batch進行子模型參數(shù)的更新和執(zhí)行訓練。KD-ATFF的算法流程如下。

輸入：訓練數(shù)據(jù)集D；訓練Epoch數(shù)；分支數(shù)n

輸出：n個訓練好的模型{θ1,θ2,θ3,…,θn}

2.whilee≤do

3.使用公式(5)計算所有分支的預測{q1,q2,…,qn}

5.計算每個子模型的輸出：(z1,z2,…,zn)

6.通過FFM獲取每個分支的權重

7.使用公式(4)計算目標logits

8.利用公式(8)計算蒸餾損失LKL

10.e=e+1

11.end while

2 實驗分析

2.1 實驗設置與數(shù)據(jù)集

實驗采用CIFAR-10和CIFAR-100多類別分類基準數(shù)據(jù)集。CIFAR-10是1個自然圖像數(shù)據(jù)集，包含從10個對象類中提取的50 000/10 000個訓練/測試樣本(總共60 000個圖像)，每個類有6 000個大小為32×32像素的圖像。CIFAR-100與CIFAR-10類似，也包含50 000/10 000個訓練/測試圖像，但覆蓋100個細粒度類，每個類別有600張圖片。實驗在CIFAR-10/100數(shù)據(jù)集上的batchsize設置為256。

實驗所使用的學生網(wǎng)絡包括：ResNet-32(3.3 M)[2]，ResNet-110(0.5 M)[2]和MobileNet(1.7 M)[10]。分支m設置為2;θ，β，T分別為1,1，3;選取隨機梯度下降，SGD為優(yōu)化器;模型的學習率初始化為0.1，并且每80周期減少為原來的1/10。采用top-1分類錯誤率，將所有模型的訓練結(jié)果取平均值。模型訓練和測試的計算成本，使用浮點運算(FLOPs)標準。

2.2 實驗結(jié)果分析

將KD-ATFF方法與幾種有代表性的蒸餾方法進行比較，采用不同的骨干網(wǎng)絡分別在CIFAR-10/100數(shù)據(jù)集上進行實驗。不同骨干網(wǎng)絡在CIFAR-10/100上的top-1錯誤率見表1。

表1 不同骨干網(wǎng)絡在CIFAR-10/100上的top-1錯誤率

由表1可知， KD-ATFF與其他幾種方法相比，top-1錯誤率有明顯降低，能適用于不同網(wǎng)絡且能得到良好的分類效果。具體來說，KD-ATFF模型在CIFAR-10上以ResNet-32或MobileNet為骨干網(wǎng)絡時，比原始Baseline的top-1錯誤率降低了約30%；在CIFAR-100上以ResNet-110為骨干網(wǎng)絡時，top-1錯誤率比DML降低了1.57%，比ONE降低了1.31%；以MobileNet為骨干網(wǎng)絡時，top-1錯誤率比DML降低了1.76%；在CIFAR-10上以ResNet-32為骨干網(wǎng)絡以及在CIFAR-100上以ResNet-110為骨干網(wǎng)絡時，與最新的OKDDip方法相比，top-1錯誤率都能夠與其比肩甚至比其更優(yōu)。實驗結(jié)果表明，KD-ATFF提高了模型的泛化性，訓練出的模型更加高效，對提升模型準確率有很大的貢獻。

2.3 消融實驗

為驗證CL模塊和特征融合模塊的有效性，在CIFAR-100數(shù)據(jù)集上使用ResNet-110為骨干網(wǎng)絡進行消融研究，將這2個模塊與ONE中的Gate模塊進行比較。消融實驗結(jié)果見表2。由表2可知，當只使用CL模塊時的性能已經(jīng)略超過其他方法。這說明CL模塊在不同維度的注意力轉(zhuǎn)移能學習到更多的知識。與ONE中的Gate模塊相比，CL的top-1錯誤率降低了0.76%。FFM與CL模塊同時工作時整體模型的改善更為明顯。與獨立的CL模塊相比，top-1錯誤率降低了1.31%。實驗結(jié)果表明，CL模塊在整體性能改善中發(fā)揮了重要的作用，特征融合模塊可以明顯增強分支間的多樣性。

表2 消融實驗結(jié)果

3 結(jié)論

結(jié)合注意力轉(zhuǎn)移與特征融合的在線知識蒸餾方法(KD-ATFF)是一種可以在不需預先訓練教師模型的前提下訓練學生模型的改進在線蒸餾模型?；谀K間的差異性，KD-ATFF引入了注意力機制，讓不同維度的模塊互相學習，在各子模型的最后1個block的特征輸出加入特征融合模塊，分配不同權重組成集成教師指導各子模型訓練，以提升整體模型的性能。與其他幾種代表性的在線知識蒸餾方法相比，KD-ATFF的top-1錯誤率明顯降低，驗證了注意力轉(zhuǎn)移和CL模塊以及特征融合模塊的有效性。