基于知識蒸餾的特征精煉相互學習方法

文 強，郭 濤，王 濤，李貴洋，鄒俊穎

(四川師范大學 計算機科學學院，四川 成都 610101)

0 引 言

神經網絡在解決很多復雜問題上的能力有了指數級的進步，已經應用于計算機視覺的各個領域[1-3]。由于訓練模型包含大量參數，部署運行需要占用較大存儲空間且執行速度緩慢，限制了其在移動設備或嵌入式設備等低內存或有快速執行要求的應用程序和平臺中的部署與運行[4，5]。知識蒸餾是一種有效加速模型速度和壓縮模型大小的方法，在模型大小和準確率上能取得很好的平衡，逐漸成為研究熱點。

知識蒸餾是將知識從預訓練的教師網絡轉移到未經訓練的學生網絡的模型壓縮方法，在部署階段，較小的學生網絡取代大型教師網絡。傳統的知識蒸餾通常采用單向知識轉移和兩階段訓練程序，即提前訓練好教師網絡，學生網絡再從教師網絡單向獲取知識。然而，復雜容量的教師模型需要較長訓練時間且占用大量的存儲。為了減少訓練大容量高性能的復雜網絡和在沒有預訓練的教師網絡等情況下，提高小型輕量型網絡的準確率以及降低模型大小，自我知識蒸餾和相互蒸餾等方法得到了快速發展。

本文提出一種基于知識蒸餾的特征精煉相互學習方法(feature refining mutual learning，FRML)。該方法在沒有教師網絡的情況下，給相互學習的學生網絡配備輔助網絡，通過對處理多尺度特征網絡的調整達到對知識提煉目的，生成精煉特征圖。輔助網絡將精煉的知識傳遞給其輔助的學生網絡，而學生網絡之間通過協作訓練，學習互相精煉后的知識，優化自身網絡，提高網絡的泛化能力以及其對視覺任務的適用性。

1 相關工作

1.1 知識蒸餾

知識蒸餾的主要思想源于Buciluǎ等提出的模型壓縮[6]，即將信息從大模型或模型集合轉移到訓練小模型，使得小模型近似大模型，從而不顯著降低準確率。知識蒸餾是一種類似于遷移學習的模型壓縮方法，但知識蒸餾不同于遷移學習。知識蒸餾的目標是提供較小的網絡模型來解決與較大網絡模型相同的任務；而遷移學習的目標是減少網絡模型的訓練時間，以解決類似于其它網絡模型所解決的任務。知識蒸餾通過改變正在訓練學生網絡的損失函數來實現其目標，以解釋預訓練教師網絡的隱藏層輸出；而遷移學習通過將預訓練網絡模型的參數遷移到新模型來實現其目標。傳統知識蒸餾傳遞知識的方法為，給定一個logitsz向量作為教師網絡的最后一個全連接層的輸出，zi是第i類的logit，采用softmax函數估計輸入屬于第i類的概率pi， 引入溫度因子T控制每個軟標簽的重要性，如式(1)

(1)

T越高在各類別上產生越弱的概率分布。通過以上方法，教師網絡獲得的軟標簽預測包含暗知識，并且可以用作監督，以將知識從教師網絡轉移到學生網絡。教師網絡提供不同的軟標簽會給同一個學生網絡帶來不同的表現。

1.1.1 自我知識蒸餾

自我知識蒸餾不同于傳統知識蒸餾，它不需要提前準備大型復雜的教師網絡，即學生網絡利用自身知識來提高訓練自身網絡的有效性。Zhang等[7]提出在只有學生網絡的情況下，將深層網絡的部分知識蒸餾到網絡較淺部分，用中間隱藏層特征對輸出進行分類。Hou等[8]提出用于通道檢測的自注意蒸餾方法，該網絡利用自身各層的注意力圖作為下層蒸餾目標，以減少模型參數量。Zhu等[9]為了降低計算成本，提出添加額外分支，使模型參數和中間隱藏層的估計特征多樣化，其中每個分支表示一個學生網絡，最后集成所有的分支網絡進行輸出。上述自我知識蒸餾方法基本上沒有對網絡進行增強，導致缺乏更精煉知識。Ji等[10]提出在理想的情況下，可以通過從復雜集成模型中提取精煉知識以改進特征圖蒸餾。

1.1.2 相互蒸餾

相互蒸餾在沒有復雜預訓練的教師網絡情況下，教師網絡和學生網絡同時學習，教師網絡會隨著學生網絡的更新而更新。Zhang等[11]提出深度相互學習，在訓練過程中，任何一個網絡均可作為學生，其它網絡作為老師，相互教學以提高模型準確率。Chen等[12]提出將輔助同伴(auxi-liary peers)和組長(group leader)引入深度相互學習中，以形成一組多樣化的同伴網絡。Kim等[13]提出采用特征融合模塊構建教師網絡，將多個并行神經網絡訓練成子網絡，利用融合模塊將子網絡知識結合起來，傳遞給教師網絡，教師網絡再將其知識傳遞給每個子網，以在線知識提取的方式相互教學。但是這種相互教學的方式，在訓練過程中會丟失局部信息，阻礙了其對各種視覺任務適用性，這種知識提煉方法，單個教師網絡或者學生網絡均接收不到精煉的特征圖。

1.2 特征網絡

特征網絡又叫特征金字塔網絡[14](feature pyramid networks，FPN)，是一種以任意尺寸的單尺度圖像為輸入，以全卷積方式在多個層次上輸出按比例大小的特征圖的特征提取器。自從FPN被提出來以后，均只有自上向下的單向融合。Liu等[15]提出路徑聚合網絡(path aggregation network，PANet)，為FPN引入額外自下而上的融合路徑，以實現網絡中層和層之間的短路徑連接，從而實現二次融合。Tan等[16]提出一種加權雙向特征金字塔網絡(weighted bi-directional feature pyramid network，BiFPN)，該網絡能簡單快速地進行多尺度特征融合。BiFPN采用與PANet相同的自上向下和自下向上網絡結構，同時增加部分跳躍連接，并將該網絡結構融合構造為一個模塊，使其可以重復堆疊，增強特征融合。同時，BiFPN為每個輸入特征增加額外權重，允許網絡學習每個輸入特征的重要性，從而獲得精煉的特征。

2 FRML模型

2.1 基本定義

定義1 數據集：訓練圖像數據集D={(x1，y1)，(x2，y2)，…，(xN，yN)}， 其中N為樣本大小，yi∈{1，2，…，M}，M表示類別大小。

定義2 學生網絡：Θ1，Θ2為學生網絡，并且Θ1=Θ2，Θ1有n個塊。設Sij為樣本i輸入Θ1的第j個塊的特征映射，cj為Θ1的第j個塊的通道維度。

定義3 輔助網絡：Θ′1、Θ′2分別為Θ1、Θ2的輔助網絡，并且Θ′1=Θ′2，Θ′1有n個塊。設Tij為輸入特征映射，為了簡化符號，在本文的其余部分省略了索引i。學生網絡與輔助網絡的特征映射塊數量相同。

定義4 logit與預測值：logitz1和z2分別為Θ1和Θ2的完全連接層的輸出；z′1和z′2分別為Θ′1和Θ′2完全連接層的輸出；p1和p2分別為Θ1和Θ2輸出預測值；p′1和p′2分別為Θ′1和Θ′2輸出預測值。

2.2 模型結構

FRML模型由兩個對稱的學生網絡Θ1和Θ2以及其輔助網絡Θ′1和Θ′2構成。首先圖像分別輸入該模型的Θ1和Θ2， 提取到n個特征圖塊Sj， 并輸出特征圖z1，z2； 其次將每個特征圖塊Sj分別輸入到Θ′1和Θ′2， 得到精煉的特征圖塊Tj， 并輸出特征圖z′1和z′2； 根據每個輸出特征圖z1，z′1，z2，z′2， 計算得到每個輸出預測值p1，p′1，p2，p′2。 最后Θ1和Θ2根據自身的特征圖塊Sj與Θ′1和Θ′2的特征圖塊Tj、 輸出預測值、真實標簽，做出相應計算，得到每個部分的損失，并計算總損失以及更新權重。兩個學生網絡相互訓練，相互學習到精煉的特征，從而提高模型的性能。FRML結構如圖1所示。

圖1 FRML結構

2.2.1 學生網絡

FRML中兩個學生網絡Θ1和Θ2采用相同的ResNet[17]網絡，該網絡包含4層block，使用每個階段的最后一個殘差結構的特征圖作為輸出，表示為 {S1，S2，S3，S4}， 通過完全連接層輸出特征圖z1和z2， 完成每個階段的特征提取，如圖1所示。

Θ1和Θ2同時訓練，相互傳遞知識，進行交流。樣本xi輸入Θ1之后所屬類別m的概率通過式(2)計算

(2)

對于多分類任務，采用交叉熵函數計算預測值與真實標簽的誤差，如式(3)、式(4)

(3)

指示函數I定義為

(4)

為提高Θ1在測試實例上的泛化性能，引入與Θ1相互學習的Θ2， 以Θ2后驗概率p2的形式提供訓練經驗。為量化Θ1和Θ2的預測值p1和p2， 使用Kullback Leibler(KL)散度。從p1到p2的KL散度計算如式(5)

(5)

至此，Θ1和Θ2相互學習損失LΘ1和LΘ2分別為式(6)、式(7)

LΘ1=LCE1+DKL(p2‖p1)

(6)

LΘ2=LCE2+DKL(p1‖p2)

(7)

兩個學生網絡通過相互學習的方式，每個學生網絡均能學習訓練實例的真實標簽，即學生網絡受監督損失LCE指導以及匹配其一起相互學習的概率估計，即KL散度損失。

2.2.2 輔助網絡

為幫助Θ1提供精煉的特征映射和軟標簽，給相互學習的Θ1添加Θ′1。 假定Θ1有n個塊，輔助網絡的輸入為學生網絡的特征映射S1，…，Sn。 輔助網絡采用來自PANet[15]自下而上的路徑和FPN[21]自上而下的路徑，以適應分類任務。輔助網絡結構如圖2所示。

圖2 輔助網絡結構

在進行自上而下的路徑之前，利用式(8)的橫向卷積層將特征映射輸入進輔助網絡

Lj=Conv(Sj；dj)

(8)

Conv為輸出維度為dj的卷積運算。dj取決于Θ1輸出特征映射的通道維度cj。

Sj經過橫向卷積得到Lj后，通過自上而下路徑，融合不同層次特征，提取低層信息Pj， 再通過自下而上路徑，進行二次特征融合，得到Tj， 如式(9)和式(10)所示

(9)

(10)

通過圖2可以看出，首先，如果節點僅包含輸入，并沒有進行特征融合，那么它對于特征網絡的貢獻是極小的。因此，將P4和P1移除，從而提高模型計算效率。其次，針對同一層級，添加一條從起始輸入到輸出的連接，即L2到T2和L3到T3， 在不增加任何參數的同時，融合更多特征。最后，添加L4到P3和P2到T1的連接，創建完整的自上而下路徑。其中，{P2，P3} 和 {T1，T2，T3，T4} 通過3×3 Conv得到。對于參數wP和wT， 使用快速歸一化融合技術[16]，引入可學習的權重，用于學習不同輸入特征的重要性。Resize運算符分別代表自上而下路徑上的上采樣和自下而上路徑上的下采樣；對上采樣使用雙向線性插值，即融合上一層2×UpSample的特征，對下采樣使用最大池化，即融合下一層3×3/2 Maxpool的特征。最后在自下而上路徑的頂部增加全連接層輸出z′1， 增加softmax層來輸出p′1預測類別。

(11)

與其它自我知識蒸餾方法類似，FRML通過軟標簽進行蒸餾，提供軟標簽蒸餾損失，如式(12)所示

(12)

其中，T是溫度系數。此外，Θ′1使用交叉熵損失函數學習真實標簽，通過式(3)計算得出L′CE1。

2.2.3 損失函數

(13)

(14)

其中，α和β為超參數，本文選擇α∈[1，2，3] 和β∈[100，200]。 對于兩個學生網絡及其輔助網絡，通過反向傳播算法進行優化。FRML模型通過交叉熵損失、KL散度損失、自特征蒸餾損失和軟標簽蒸餾損失更新參數。

2.3 算法流程

精煉特征相互學習方法分為兩個階段：首先，兩個學生網絡通過輔助網絡分別獲取到精煉的知識，然后，兩個學生網絡相互學習精煉的特征，兩個學生網絡互相微調，以提高自身的泛化能力。FRML模型的算法的偽代碼如算法1所示。

算法1：特征精煉相互學習算法

輸入：學生網絡Θ1和Θ2； 輔助網絡Θ′1和Θ′2； 訓練數據集；學習率、小批量尺寸以及α，β

輸出：更高分類準確率的學生網絡

(1)隨機初始化Θ1和Θ2中所有網絡層的參數

(2)for 每個小批量

2.1 輸入樣本xi, 真實標簽yi;

2.2 輸入Θ1, 將Sj輸入到Θ′1, 得到Tj并得到特征圖z1,z′1;

2.3 輸入Θ2, 將Sj輸入到Θ′2, 得到Tj并得到特征圖z2,z′2;

2.4 特征圖通過式(2)分別得到預測值p1、p′1和p2、p′2;

2.5 通過式(3)計算交叉熵損失LCE1, L′CE1和LCE2, L′CE2;

2.10 計算Θ1的總損失

2.11 使用反向傳播算法更新Θ1的權重;

2.12 計算Θ2的總損失

2.13 使用反向傳播算法更新Θ2的權重;

(3)end for

3 實驗結果及分析

3.1 實驗數據集

本文使用了6種數據集，其中CIFAR100和TINY由小比例圖像組成，這里將TINY圖像的尺寸調整為與CIFAR100圖像尺寸相同，即32×32大小。CUB200、MIT67、Stanford40和Dogs均為細粒度視覺分類(fine-grained visual categorization，FGVC)任務的圖像數據集。實驗使用數據集概況見表1。

表1 本文所用數據集概況

3.2 實驗設置

FRML模型在數據集CIFAR100和TINY上采用相同的ResNet18[17]和WRN-16-2[20]作為學生網絡。為了使ResNet18適應小型數據集，將ResNet18的第一個卷積層修改為3×3的內核大小，采用單個步長和單個填充，并刪除最大池化操作。在FGVC任務數據集上采用標準的ResNet18作為學生網絡。

實驗基于深度學習框架Pytorch 1.8.0、Python 3.8.8 實現，并在NVIDIA GeForce RTX 3090 GPU上完成。

對于所有訓練樣本，使用隨機裁剪和翻轉來增加數據的多樣性。在訓練過程中，采用隨機梯度下降(SGD)方法，初始的學習率為0.099，動量參數為0.9，優化器中權重衰減參數為0.0001；溫度系數T為4；通道寬度參數w為2；超參數[10]的設置見表2。

表2 超參數設置

3.3 分類性能對比實驗

對于圖像分類任務，采用測試集Top-1準確率(Accuracy)作為評價指標。本實驗完成了FRML模型與Baseline、DML[11]、FRSKD[10]這3種分類方法在CIFAR100、TINY和FGVC任務數據集上的準確率比較。實驗結果見表3和表4。

表3 FRML在CIFAR100和TINY數據集上的分類準確率

表4 FRML在FGVC任務數據集上的分類準確率

從表3可知，FRSKD和DML均提高了標準分類器網絡的性能，但FRML表現出更好的性能，與基線FRSKD相比，平均準確率提高了1.49%，與基線DML相比，平均準確率提高了2.30%。

表4顯示了在FGVC任務數據集上的分類準確率。在4種數據集上的實驗結果可知，FRML表現出比其它方法更好的性能。與基線FRSKD相比，平均準確率提高了1.85%，與基線DML相比，平均準確率提高了8.08%。通過實驗驗證，FRML提高了相互學習在細粒度視覺任務的準確率。

3.4 知識蒸餾分類實驗

在知識蒸餾分類實驗中，將FRML與FitNet[21]、ATT[18]、Overhaul[19]等知識蒸餾方法做對比實驗。教師網絡設置為預訓練的ResNet34，學生網絡設置為未經訓練的ResNet18。其中，分類準確率的最大值用粗體標識，實驗結果見表5。

表5 知識蒸餾上分類準確率

從表5中可知，來自強大的教師網絡的傳統知識蒸餾方法雖然提高了學生的表現。但FRML在所有數據集上優于使用教師網絡的知識蒸餾方法，這意味著在相互學習的過程中，通過精煉特征傳遞給學生網絡，比預訓練的教師網絡更好。

3.5 擴展不同學生網絡的實驗

為了進一步提高模型的適用性，不限制模型架構，本實驗將兩個學生網絡擴展成不同結構的網絡。其中，在數據集CIFAR100和TINY上分別采用WRN-16-2和ResNet18作為學生網絡Θ1和Θ2， 在FGVC任務數據集上分別采用ResNet18和ResNet32作為學生網絡Θ1和Θ2。 FRML模型在不同網絡架構以及不同數據集上分類準確率見表6。

表6 不同學生網絡分類準確率

從表6可知，相比較獨立訓練的網絡和DML，基于不同網絡的組合，FRML均提高了每個學生網絡的性能。與DML相比，在數據集CIFAR100和TINY上，WRN-16-2網絡的平均準確率提高了1.65%，ResNet18網絡的平均準確率提高了2.00%。在FGVC任務數據集上采用ResNet18網絡的平均準確率提高了3.84%，采用ResNet32網絡的平均準確率提高了4.03%。由此可知，不同網絡架構或不同大小的網絡，通過特征精煉相互學習后，性能均得到提高。實驗結果表明，在不同網絡架構上，本方法同樣能夠改進相互學習的不足，提高不同網絡的泛化性能。

3.6 模型收斂性實驗

本實驗為FRML模型在FGVC任務數據集上訓練次數以及測試分類準確率上的收斂實驗，結果如圖3所示。圖3(a)～圖3(d)分別為模型在CUB200、MIT67、Dogs、Stanford40數據集上的收斂效果。

圖3 FRML在FGVC任務數據集上的收斂曲線

從圖3(a)～圖3(d)可以看出，FRML模型在FGVC任務數據集上均具有良好的收斂性。當訓練次數在100次之前，準確率的提升速度較快；在訓練到100次左右時有較大提升；在之后的訓練次數中，分類準確率小幅度振動，模型開始整體進入收斂狀態；訓練到150次左右時，可以看出準確率有個小幅度的上升；在150次之后，分類準確率波動趨于穩定，模型基本達到收斂。

此外，從圖3(a)～圖3(d)可以看出，對比FRML和DML，FRML在FGVC任務數據集上的分類準確率明顯高于DML。實驗結果表明本方法解決了相互學習上的不足，提高了學生網絡在相互學習的泛化性能，以及其對視覺任務的適用性。

4 結束語

本文在沒有預訓練的教師網絡的情況下，提出一種基于知識蒸餾的特征精煉相互學習方法，解決兩個學生網絡在相互學習時出現圖像局部信息丟失，而造成單個學生網絡接收精煉特征圖不完整的問題。通過對每個學生網絡配備特征提取輔助網絡進行特征融合，利用輔助網絡輸出精煉特征映射和軟標簽對學生網絡進行訓練，并保留圖像特征的局部信息，為學生網絡之間傳遞精華知識，以達到最優效果。通過改進相互學習策略，提高了視覺任務的適用性。后續的研究是將學生網絡的數量擴展到3個及以上以進一步提高其性能。