基于關系型蒸餾的分步神經網絡壓縮方法①

劉 昊, 張曉濱

(西安工程大學 計算機學院, 西安 710048)

近些年來, 深度神經網絡獲得了越來越多的關注,且在各種應用領域都取得了很好的成果, 例如計算機視覺[1,2]、自然語言處理[3]等. 但是深層的神經網絡往往受限于其模型大和復雜度高的特性, 很難直接部署到計算和存儲能力有限的設備, 例如移動設備和嵌入式傳感器等[4]. 因此, 在保證一定的準確度和精度的條件下, 合理的對深度神經網絡進行模型壓縮和優化成為了一個很有研究價值的問題. 知識蒸餾就是模型壓縮方法中的一種非常典型的方法.

知識蒸餾的概念最早由Bucila等人[5]在2006年提出, 由Hindon等人[6]在2015年重拾并普及. 該方法的核心思想在于, 將復雜的、學習能力強的教師網絡學到的知識遷移到較小的學生網絡, 從而提升學生網絡的精度. 但是在原始的知識蒸餾方法中, 學生網絡只是學習了教師網絡模型的輸出, 并沒有考慮教師網絡中的其他信息. 隨后Romero等人[7]很快提出在知識蒸餾方法基礎上, 引導學生模型先學習教師網絡隱含層的特征表達來作為一個預訓練模型, 再通過知識蒸餾得到目標學生網絡, 以使得學生網絡可以學習到更豐富的信息. Zagoruyko等人[8]在2016年提出使用教師網絡中的注意力圖進一步訓練學生網絡, 通過定義注意力機制, 讓學生網絡擬合教師網絡的注意力映射, 從而提升學生網絡的效果, 達到知識蒸餾的目的. Yim等人[9]在2017年提出讓學生網絡去學習教師網絡中層與層之間的關系而不是輸出, 層間關系是用關系矩陣來定義的, 相當于讓學生網絡擬合了教師網絡的學習過程, 這個思路也得到了不錯的蒸餾效果. Chen 等人[10]提出將原始蒸餾和文獻[7]中提出的隱層蒸餾結合起來針對物體檢測任務做知識蒸餾. Huang等人[11]提出將知識蒸餾的遷移過程看作是特征分布擬合過程, 并采用域適應中常用的最大平均差異來優化. Chen等人[12]提出使用跨樣本相似性作為新的知識來讓學生網絡進行學習, 在人員檢測和圖像檢索任務中都收到了不錯的效果. Heo等人[13]提出遷移兩種知識, 包括ReLU之后的特征響應大小以及每一個神經元的激活狀態讓學生網絡進行學習. Heo等人[14]在AAAI上也提出了基于激活邊界的知識蒸餾. Saputra等人[15]的工作對用于回歸任務的網絡進行知識蒸餾有一定的實踐指導價值. Mishra等人[16]利用知識蒸餾技術提升低精度網絡的分類性能. Wonpyo等人[17]在2019年提出關系型知識蒸餾方法, 讓學生網絡學習由教師網絡的多個輸出組成的結構性信息, 包括二元的距離結構信息和三元的角度結構信息, 在蒸餾的效果上相較之前的方法有了更明顯的提升.

上述方法都是在教師網絡知識遷移類型上來進行優化, 并沒有考慮到教師網絡和學生網絡層數差距過大時, 學生網絡的擬合能力受到限制從而導致蒸餾效果下降的問題. 對于這個問題, Cho等人[18]提出提前終止教師網絡的訓練來使學生網絡更好的擬合教師網絡,這個方法在一定程度上降低了師生差距過大導致的影響, 但也損失了部分教師網絡的信息. Jin等人[19]的工作受課程學習啟發, 提出路由約束學習, 讓學生網絡學習教師網絡訓練過程中不同階段的狀態, 由易到難不斷學習和優化. Mirzadeh等人[20]于2019年引出了“助教”的概念, 即在教師網絡和學生網絡中間使用一個中間的蒸餾網絡進行過渡, 一定程度上避免了教師網絡和學生網絡層數差距過大的問題. 但是這個方法僅考慮了原始的知識蒸餾方法的優化, 蒸餾得到的學生模型效果并不是十分理想.

本文針對教師網絡和學生網絡層數差距過大影響蒸餾效果的問題, 基于關系型知識蒸餾方法, 提出了基于關系型蒸餾的分步神經網絡壓縮方法. 首先訓練深層的神經網絡作為教師網絡, 選取中間網絡來學習教師網絡的單體輸出與關系型輸出作為過渡, 隨后讓目標學生網絡學習中間網絡的單體輸出與關系型知識.通過分步的關系型蒸餾來緩解教師網絡和學生網絡的層數差距, 同時在每一次關系型蒸餾過程中增加單體輸出來豐富學生網絡的信息, 實現高層到低層的神經網絡壓縮.

1 關系型分步知識蒸餾

本文以關系型知識蒸餾作為基礎, 通過分步的蒸餾方法有效的避免了教師網絡和學生網絡層數差距過大時蒸餾效果下降的問題, 以使得較低層數的學生網絡也可以很好的學習到較高層數的教師網絡的知識,從而獲得更高質量的學生網絡.

1.1 關系型蒸餾

傳統的知識蒸餾方法通常只考慮到教師網絡的輸出表現或在此基礎上的改進, 很少考慮到教師網絡的結構信息. 相比之下, 關系型蒸餾引入了教師模型的多個輸出組成結構單元來讓學生網絡進行學習. 由于這些信息更能體現出教師模型的結構化特征, 關系型蒸餾使得學生模型可以得到更好的指導.

關系型蒸餾的整體流程如圖1.

圖1 關系型蒸餾流程

步驟1. 確定教師網絡和目標學生網絡的整體結構, 并用正常的訓練方式訓練教師網絡；

步驟2. 以教師網絡的多個輸出組成的結構性知識作為監督來訓練學生網絡；

步驟3. 訓練完成的學生網絡即為目標模型.

其中,t1,···,tn代表教師模型的一個批次中的多個輸出,s1,···,sn代表一個批次中學生模型的多個輸出,關系矩陣rt和rs分別為教師和學生模型的輸出經過距離/角度關系函數變換得到的關系矩陣.

從圖1中可以看出, 模型整體是通過使用教師網絡中距離或角度信息的關系矩陣作為監督來訓練學生模型從而實現知識蒸餾, 因此, 損失函數可以定義如下:

其中,x1,···,xn代 表一個批次樣本中的n元組,t1,···,tn代表教師模型的多個輸出,s1,···,sn代表學生模型的多個輸出, φ為給定n元組的關系函數,l是教師模型的結構信息與學生模型的結構信息的損失函數. 考慮到效率和運算成本, φ關系函數的選取目前只給出兩種情況, 一種是距離關系函數φd, 此時n取值為2； 一種是角度關系函數φa, 此時n取值為3.

當選取距離關系函數時, 將每個批次的m個樣本分成二元組 (ti,tj)分別計算歐幾里得距離, 距離計算方法如下所示:

其中, μ是距離的正則化參數, 為了更好的關聯其他二元組的距離數據, μ定義如下:

通過兩兩計算距離, 可以得到一個m×m的距離矩陣作為關系型的距離結構信息, 矩陣中包含了批次中每一個樣本到其他所有樣本的距離關系. 學生模型通過學習這個關系型距離結構信息來實現蒸餾學習, 此時損失函數如下:

當選取角度關系函數時, 將每個批次的m個樣本分成若干三元組分別計算角度, 角度計算方法如下所示:

通過每一個三元組計算角度, 得到一個m×m×m的角度矩陣作為關系型的角度結構信息, 由于角度信息是一個更高階的屬性, 因此可以更高效的傳遞關系信息. 學生模型通過學習這個關系型角度結構信息來實現蒸餾學習, 此時損失函數如下:

距離結構信息和角度結構信息可以通過給定權重,共同作為監督來訓練學生網絡, 讓學生網絡學習到更豐富的結構型信息. 但是如上一小節所述, 關系型蒸餾方法同樣存在教師網絡和學生網絡層數差距過大時損失蒸餾效果的問題, 因此本文引入了分步蒸餾的思想來進行進一步的改進.

1.2 基于關系型蒸餾的分步神經網絡壓縮方法

本文方法以關系型蒸餾方法為基礎, 在教師網絡和學生網絡中間增加了一個中間網絡作為過渡, 分步進行知識蒸餾, 同時在每一步的蒸餾過程中都學習了額外的單體信息. 單體信息是由教師或學生的輸出首先除以一個軟化系數t[5], 再經過Softmax變換得到的,軟化系數t用來緩和教師網絡的原輸出, 取值越大, 輸出的分布越緩和.

模型的整體訓練流程如算法1.

算法1. 模型訓練流程x1,···,xn 輸入: 多批次的圖片樣本輸出: 學生網絡模型S x1,···,xn t1,···,tn步驟1. 確定教師網絡和學生網絡的模型結構, 確定中間網絡的層數及結構, 隨后使用樣本 訓練教師網絡T, 得到教師網絡的輸出信息a1,···,an步驟2. 根據式(2)和式(5)計算得到教師網絡的二元距離關系矩陣rdt和三元角度關系矩陣rat, 與教師網絡軟化的單體輸出softt協同作為監督訓練中間網絡A, 得到中間網絡的輸出信息步驟3. 根據式(2)和式(5)計算得到中間網絡的二元距離關系矩陣rda和三元角度關系矩陣raa, 與中間網絡軟化的單體輸出softa協同作為監督訓練學生網絡S步驟4. 訓練好的學生網絡S即為最終的目標模型

如圖2所示, 模型整體經過了教師到中間網絡和中間到學生網絡的蒸餾過程. 以教師到中間網絡為例,中間網絡模型在訓練過程中會學習并擬合教師模型的3種知識, 包括教師模型軟化的單個輸出集合softt； 教師模型輸出的二元距離值組成的距離關系矩陣rdt； 教師模型輸出的三元角度值組成的角度關系矩陣rat.

圖2 本文方法模型框架圖

圖2中, 軟化輸出通過交叉熵來計算損失值； 二元距離關系矩陣由教師網絡的輸出通過式(2)計算得到, 三元角度關系矩陣由教師網絡的輸出通過式(5)計算得到, 且二元距離矩陣和三元角度矩陣均使用了Smooth L1損失來計算損失值.

顯然, 教師到中間網絡蒸餾過程的整體的損失函數由3部分組成, 分別是單體損失、距離關系損失和角度關系損失, 具體可以表示為:

其中,Lzmrkd為教師向中間網絡蒸餾的整體損失函數,Ltkd為中間網絡的單體輸出損失,Ltrkd-d和Ltrkd-a分別為2.1節中提到的中間網絡的距離關系損失和角度關系損失, λkd為 單體輸出權重, λd為距離結構信息權重, λa為角度結構信息權重.

經過教師模型到中間網絡的知識蒸餾, 可以得到一個蒸餾了教師網絡中各種知識的較淺的中間網絡,該中間網絡的輸出將作為監督進一步訓練學生網絡.與中間網絡訓練過程類似, 學生網絡模型在訓練過程中也會擬合中間網絡的3種知識, 包括中間網絡軟化的單個輸出集合； 中間網絡輸出的二元距離值組成的距離關系矩陣； 中間網絡輸出的三元角度值組成的角度關系矩陣.

中間網絡到學生網絡蒸餾過程的整體損失函數具體可表示為:

其中,Lsmrkd為中間網絡向學生網絡蒸餾的整體損失函數,Lskd為學生網絡的單體輸出損失,Lsrkd-d和Lsrkd-a分別為學生網絡的距離關系損失和角度關系損失.

模型整體通過增加中間網絡分步進行知識蒸餾來緩和教師網絡和學生網絡之間的層數差距, 同時本文在每一次蒸餾過程中都使用了教師網絡的3種知識協同對學生網絡進行指導, 學生網絡不僅可以遷移獲得兩種關系型信息, 還可以學習到教師網絡單體輸出帶來的直接信息, 從而使得學生網絡可以更好的擬合教師網絡的學習能力.

2 實驗分析

本文實驗所采用的數據集為CIFAR-10和CIFAR-100圖像分類數據集, 其中CIFAR-10數據集包括60 000張尺寸為32×32的彩色圖像, 圖像分為10類, 每一個分類中都有6000個圖像, 包括5000個訓練圖像和1000個測試圖像； CIFAR-100和CIFAR-10類似, 但是有100個類, 每個類包含600個圖像, 包括500個訓練圖像和100個測試圖像. 為了更直觀的體現深度變化,采用ResNet作為實驗模型的主體, 其中ResNet-110作為教師網絡模型, ResNet-20作為目標學生網絡模型,ResNet-44作為中間網絡. 實驗同時將本文的方法與原始的知識蒸餾KD (Knowledge Distillation)模型、文獻[7]提到的FT (FitNet)模型、文獻[8]提到的AT(Attention Transfer)模型、文獻[9]提到的FSP (Fast optimization, network minimization and transfer learning)模型、文獻[11]提到的NST (Neural Selective Transfer)模型、文獻[17]提到的RKD(Relational Knowledge Distillation)模型以及文獻[20]提到的TAKD (Knowledge Distillation via Teacher Assistant)模型在相同參數和模型設置下做了對比實驗.

2.1 實驗設置

本文的模型均基于PyTorch實現, 初始學習率設定為0.1, batch-size設定為64, epochs設定為100. 蒸餾過程中的參數中, KD模型軟化系數T為4.0, 權重λkd為16, AT模型權重設定為2.0, RKD模型中距離結構信息權重 λd為25.0, 角度結構信息權重λa為50.0.ResNet-20、ResNet-110、ResNet-44的網絡結構參照文獻[21]中實現. 受限于機器性能, 本文給出的實驗數據不代表模型最佳表現, 這里只給出實驗對比結果來進行模型的對照.

2.2 實驗結果與分析

實驗采用的主體評價指標為圖像分類準確率, 同時也給出了一部分模型參數量和召回率的實驗數據對比. 為了達到更好的對照效果, 模型共用的參數如軟化系數T, 均采用相同的取值.

首先, 本文首先分別用CIFAR-10和CIFAR-100數據集訓練教師網絡ResNet110和學生網絡ResNet20,以獲得未經過蒸餾的模型的準確率, 隨后使用本文的方法, 訓練得到蒸餾過后的ResNet20作為目標學生模型. 所得結果如表1所示.

從表1可以看出, 本文的方法得出的模型結構與原始ResNet20相同, 但是在兩個數據集上的準確率均有很大的提升.

同時列出ResNet110、ResNet44、ResNet20及本文模型的網絡層數及參數量, 如表2所示.

從表2可以看到, 本文中目標學生網絡的參數量僅有0.27 M, 遠小于教師網絡甚至是中間網絡的參數量, 且通過表1的數據來看, 本文模型相對于同結構的ResNet20, 準確率有相當大程度的提升.

表1 初始模型與本文模型準確率(%)

表2 網絡層數及參數量對比

為了說明在教師網絡和學生網絡的差距過大時加入中間網絡的影響, 以及在分步蒸餾過程中加入單體輸出信息協同結構信息來進一步監督的影響, 本文在學生網絡設定為ResNet20的情況下, 分別進行了教師網絡為ResNet20、ResNet50以及ResNet110的關系型蒸餾實驗, 同時使用TAKD的思路進行了關系型蒸餾的實驗, 所得結果如表3所示.

表3 本文模型與低層教師關系型蒸餾準確率(%)

從表3可以看出, 當教師模型分別設定為ResNet20和ResNet50時, 隨著教師網絡自己的模型效果提升,學生模型的準確率也有了一定的提升； 當教師模型設定為ResNet110時, 教師模型準確率有很大提升, 學生模型的準確率相比教師模型為ResNet50卻有下降, 甚至跟教師模型為ResNet20時的準確率相仿, 因此在教師網絡和學生網絡差距過大時實際上完全失去了知識蒸餾的意義. 使用TAKD的思路在ResNet110和ResNet20中加入中間網絡作為緩沖后, 學生模型準確率有了一定的提升. 本文模型在兩次蒸餾過程中加入單體輸出作為監督后, 可以在圖像分類中獲得更高的準確率, 達到更好的效果.

本文還與其他文獻中給出的知識蒸餾方法的準確率進行了對比, 同時增加了CIFAR-10數據集的召回率對比. 使用的教師模型均為ResNet110, 學生模型為ResNet20. 受表格格式限制, 使用C-10代表CIFAR-10數據集, C-100代表CIFAR-100數據集, acc代表準確率數據, rec代表召回率數據. 所得結果如表4所示.

表4 本文模型與其他蒸餾模型準確率及召回率(%)

從表4可以看出, 本文模型在CIFAR-10數據集上相較于其他主流模型, 準確度和召回率均有明顯提升,CIFAR-100數據集中, 模型表現跟FT模型相比雖有些許差距, 但總體來說也收到了不錯的蒸餾效果. 特別是相較于關系型蒸餾RKD模型, 本文模型在兩個數據集上均有0.2%左右的準確度提升.

以上實驗表明, 當教師網絡和學生網絡層數差距過大時, 通過選取中間網絡分步對教師模型進行包括單體輸出和結構型輸出的蒸餾, 可以在讓學生網絡學到更多關于教師網絡的信息, 遷移更豐富的知識, 進而緩解教師模型和學生模型差距過大時蒸餾效果變差的問題, 獲得相比于其他大部分主流蒸餾模型更好的效果.

3 結論與展望

本文針對關系型知識蒸餾方法中教師模型與學生模型的差距過大時蒸餾效果下降的問題, 選取中間網絡, 分步對教師網絡模型進行關系型蒸餾, 同時在每一次蒸餾過程中都在距離結構信息和角度結構信息這兩種關系型信息之外, 額外遷移了單個輸出信息來豐富學生網絡的學習. 實驗結果表明, 本文的模型在圖像分類數據集上的表現相對于原始的關系型模型以及大部分前沿的知識蒸餾模型, 分類準確率更高, 效果更好.但本文方法在中間網絡的選取以及知識遷移方式上仍有改進空間, 將在后續的研究過程進一步深化和推進.