周兆京，王曉茹，姜竹青，門愛東，馬 龍

(1.北京郵電大學(xué) 人工智能學(xué)院，北京 100876；2.北京市經(jīng)濟管理學(xué)校 信息技術(shù)系，北京 100089；3.中國人民解放軍96962部隊，北京 102206)

0 引 言

近年來，圖像超分辨率已成為一項重要的研究課題，它在目標(biāo)檢測、人臉識別和信息取證等方面有著重要的應(yīng)用價值。超分辨率旨在學(xué)習(xí)一種非線性映射，從低分辨率圖像中重建出高分辨率圖像。傳統(tǒng)插值算法主要是基于領(lǐng)域像素點進行加權(quán)計算來生成高分辨率圖像，但僅能獲得原圖本身像素點的領(lǐng)域信息，無法生成新的高頻信息，導(dǎo)致計算而得的高分辨率圖像往往不夠清晰。

隨著深度學(xué)習(xí)的快速發(fā)展，研究者們提出了許多基于卷積神經(jīng)網(wǎng)絡(luò)的超分辨率方法[1-4]。SRCNN[2]是首個將卷積神經(jīng)網(wǎng)絡(luò)引入超分辨率領(lǐng)域的算法，取得了遠(yuǎn)勝傳統(tǒng)算法的性能表現(xiàn)。SRCNN通過學(xué)習(xí)低分辨率圖像和高分辨率圖像的映射關(guān)系，可以重建出低分辨率圖像中缺失的高頻分量。增加卷積神經(jīng)網(wǎng)絡(luò)深度可以設(shè)計出卓有成效的超分辨率模型，進一步提高超分辨率網(wǎng)絡(luò)的重建效果，但是其計算復(fù)雜度和內(nèi)存占用量也急劇提升，直接在計算資源有限的設(shè)備端(如電視或手機)上實現(xiàn)它們是一個巨大的挑戰(zhàn)。

為了解決這些難題，研究者們投入大量精力去研究如何壓縮與加速神經(jīng)網(wǎng)絡(luò)模型[5]。大量的研究工作集中在網(wǎng)絡(luò)剪枝、網(wǎng)絡(luò)量化(將網(wǎng)絡(luò)權(quán)值從浮點數(shù)量化為8比特整型值)、輕量化結(jié)構(gòu)設(shè)計和知識蒸餾4個方面去優(yōu)化神經(jīng)網(wǎng)絡(luò)的推理速度。其中，知識蒸餾是一種新穎的模型壓縮方法，它通過將效果卓越的大網(wǎng)絡(luò)中的知識傳遞到一個層數(shù)更淺、復(fù)雜度更低的小網(wǎng)絡(luò)中去減輕計算負(fù)擔(dān)，無需改變網(wǎng)絡(luò)原本的結(jié)構(gòu)特點。通過傳遞大網(wǎng)絡(luò)的高級特征表達(dá)，小網(wǎng)絡(luò)能夠接收到比數(shù)據(jù)集里的標(biāo)簽更強有力的監(jiān)督信息。

本文主要研究了知識蒸餾在超分辨率網(wǎng)絡(luò)上的運用，通過大量實驗論述了目前知識蒸餾方法在超分辨率上的局限性，并提出了一種新的知識蒸餾訓(xùn)練框架，基于編解碼器的結(jié)構(gòu)提取高分辨率圖像中的先驗知識，再將其傳遞給小網(wǎng)絡(luò)，從而有效提升小網(wǎng)絡(luò)的超分辨率重建效果，實現(xiàn)超分辨率網(wǎng)絡(luò)的輕量化。

1 知識蒸餾方法

知識蒸餾[6]的主要思想是引導(dǎo)性能較弱、模型體積小的網(wǎng)絡(luò)模仿性能更優(yōu)、模型體積大的網(wǎng)絡(luò)，以獲得更優(yōu)越的性能表現(xiàn)。一般而言，在知識蒸餾的訓(xùn)練模式中，網(wǎng)絡(luò)層數(shù)深、復(fù)雜度高且性能更優(yōu)的大模型被定義為大網(wǎng)絡(luò)，網(wǎng)絡(luò)層數(shù)淺、復(fù)雜度低且性能平平的小模型被定義為小網(wǎng)絡(luò)。其關(guān)鍵在于如何定義大網(wǎng)絡(luò)中的知識，以及如何傳遞知識。目前，知識蒸餾方法按照知識的定義可大致分為3類：基于softmax層輸出的知識蒸餾[6]、基于特征的知識蒸餾[7-9]和基于特征關(guān)系的知識蒸餾[10-11]。這些方法在許多視覺任務(wù)上都取得了不錯的實驗結(jié)果。

然而，大多數(shù)知識蒸餾方法都是面向高級視覺任務(wù)而言的，如圖像分類、目標(biāo)檢測等，而鮮有針對低級視覺任務(wù)的研究。為了實現(xiàn)超分辨率網(wǎng)絡(luò)的壓縮，本文將經(jīng)典的知識蒸餾方法直接運用至超分辨率網(wǎng)絡(luò)上，對比小網(wǎng)絡(luò)采用蒸餾訓(xùn)練前后的重建效果，判斷其方法的有效性，詳細(xì)結(jié)果見表1。

表1 知識蒸餾在超分辨率網(wǎng)絡(luò)的應(yīng)用Tab.1 Application of knowledge distillation in super-resolution

表1采用文獻[14]的EDSR作為基準(zhǔn)網(wǎng)絡(luò)，大網(wǎng)絡(luò)采用與EDSR原論文相同的設(shè)置(包含32個殘差模塊、256個通道數(shù))，小網(wǎng)絡(luò)采用結(jié)構(gòu)規(guī)模更小的EDSR(包含4個殘差模塊、64個通道數(shù))，在Set5和Set14這兩個標(biāo)準(zhǔn)數(shù)據(jù)集上進行了4倍超分辨率測試，使用峰值信噪比(point signal to noise ratio, PSNR)衡量重建圖像的質(zhì)量。PSNR是圖像復(fù)原任務(wù)中使用最普遍的圖像質(zhì)量評估方法，其值越大，說明生成圖像與真實圖像越相近，圖像質(zhì)量越好。Soft-target[6]是基于softmax輸出的知識蒸餾方法，其超分辨率任務(wù)輸出的是高分辨率圖像而并無softmax輸出，實際實驗中則修改為用大網(wǎng)絡(luò)生成的圖像監(jiān)督小網(wǎng)絡(luò)。AT[8]和FitNet[9]都是典型的基于特征的知識蒸餾方法，F(xiàn)SP[10]、CCKD[12]和SPKD[13]則是基于特征關(guān)系知識蒸餾的代表。

觀察表1可得，這些知識蒸餾方法運用到超分辨率任務(wù)上，大多數(shù)都難以起到提升小網(wǎng)絡(luò)性能的積極作用，只有基于輸出圖像的soft-target方法能帶來一些改進。不同于分類等高級視覺任務(wù)，在超分辨率等低級視覺任務(wù)中，像素間的局部和全局關(guān)系尤為重要?？紤]到這一特點，諸如FSP和AT等基于特征關(guān)系的方法是從特征圖中抽象出更高級的特征表達(dá)知識，可能改變了圖像本身的空間信息，從而影響了模型恢復(fù)圖像的能力。而對于soft-target提供的加成，這可能是因為學(xué)習(xí)大網(wǎng)絡(luò)的輸出圖像比直接學(xué)習(xí)真實圖像更簡單，一定程度上降低了模型的訓(xùn)練難度。圖1展現(xiàn)了FSP與基準(zhǔn)網(wǎng)絡(luò)的訓(xùn)練過程對比情況，橫軸是訓(xùn)練批量的迭代次數(shù)，豎軸是在Set5數(shù)據(jù)集上測試的PSNR結(jié)果，baseline是未經(jīng)過知識蒸餾訓(xùn)練的小網(wǎng)絡(luò)，F(xiàn)SP是采用FSP蒸餾訓(xùn)練的小網(wǎng)絡(luò)。觀察圖1可得，知識蒸餾在訓(xùn)練前期能為小網(wǎng)絡(luò)提供一個良好的助力，加快其收斂速度，但收斂后和基準(zhǔn)網(wǎng)絡(luò)逐漸趨于一致。

圖1 FSP與基準(zhǔn)網(wǎng)絡(luò)的訓(xùn)練過程對比Fig.1 Comparison of the training process between FSP and baseline

綜上可得，目前的知識蒸餾方法直接運用在超分辨率任務(wù)上效果甚微，僅能在訓(xùn)練前期起到加速效果。為了分析知識蒸餾在超分辨率等圖像回歸任務(wù)上效果不佳的原因，本文將超分辨率網(wǎng)絡(luò)EDSR的中間層特征圖進行了可視化，見圖2。圖2中，左圖是輸入圖像，右圖是EDSR首個殘差塊輸出的特征圖。觀察EDSR的特征圖可以發(fā)現(xiàn)，該特征圖與原圖輪廓保持著高度一致，這是由超分辨率任務(wù)本身特性決定的。分類網(wǎng)絡(luò)的特征圖往往趨向于關(guān)注目標(biāo)的局部特征，其稀疏性較高，而超分辨率網(wǎng)絡(luò)提取的特征仍然與輸入圖像趨于一致，每一部分細(xì)節(jié)都影響著成像質(zhì)量，這無疑增大了通過蒸餾傳遞特征知識的難度。由此可知，要想知識蒸餾在超分辨率上發(fā)揮作用，關(guān)鍵在于大網(wǎng)絡(luò)能提供更加有益于小網(wǎng)絡(luò)訓(xùn)練的特征知識。

圖2 超分辨率網(wǎng)絡(luò)EDSR的特征圖Fig.2 Feature map of EDSR

受文獻[10]的啟發(fā)，本文針對超分辨率任務(wù)提出了一種新的知識蒸餾訓(xùn)練框架，借助編解碼器的結(jié)構(gòu)，保留真實圖像中的高頻信息，為特征知識提供更精準(zhǔn)、有用的特征，助力小網(wǎng)絡(luò)的訓(xùn)練。

2 基于編解碼器的特征知識蒸餾

2.1 構(gòu)建大網(wǎng)絡(luò)

本文的知識蒸餾訓(xùn)練框架如圖3所示。圖3中，HR表示訓(xùn)練集中的高分辨率圖像，LR表示高分辨率圖像對應(yīng)的低分辨率圖像。

圖3 基于編解碼器的特征知識蒸餾Fig.3 Feature knowledge distillation based on codec

不同于傳統(tǒng)知識蒸餾中直接從已訓(xùn)練好的大網(wǎng)絡(luò)中提取知識，本文知識蒸餾訓(xùn)練框架對大網(wǎng)絡(luò)增加了一些結(jié)構(gòu)上的訓(xùn)練約束，并使用高分辨圖像作為大網(wǎng)絡(luò)的輸入對其進行訓(xùn)練。大網(wǎng)絡(luò)分為編碼器和解碼器兩部分，編碼器對輸入的高分辨率圖像進行壓縮編碼，它將輸入圖像投影到一個低維特征空間中，生成更緊湊的特征，然后再輸入解碼器中重構(gòu)出高分辨率圖像，使大網(wǎng)絡(luò)能為超分辨率任務(wù)提取更好的特征表示。

(1)

解碼器的損失函數(shù)LD具體定義為

(2)

(2)式中，λ為退化損失的平衡系數(shù)；H、W分別表示高分辨率圖像的高和寬。(2)式由2部分組成，第1部分是生成圖像與真實圖像的MAE損失，第2部分是生成圖像退化后與解碼器輸入的MAE損失。

編碼器、解碼器共同進行訓(xùn)練，大網(wǎng)絡(luò)能從高分辨率圖像中提取高頻信息，從而在網(wǎng)絡(luò)中生成精確的特征知識，整個大網(wǎng)絡(luò)的總的損失函數(shù)LT具體定義為

LT=βTLE+LD

(3)

(3)式中，βT是編碼器損失函數(shù)的平衡系數(shù)。

2.2 構(gòu)建小網(wǎng)絡(luò)

小網(wǎng)絡(luò)采用與大網(wǎng)絡(luò)中解碼器相同的結(jié)構(gòu)，只是采用低分辨率圖像作為輸入，并使用解碼器的網(wǎng)絡(luò)權(quán)重對其進行初始化，為小網(wǎng)絡(luò)提供一個良好的訓(xùn)練起點。雖然小網(wǎng)絡(luò)和解碼器具有相同的初始參數(shù)，但由于其輸入不同，兩者提取到的特征也大相徑庭。

不同于解碼器，小網(wǎng)絡(luò)的損失函數(shù)不僅包含生成圖像的重建損失以及退化損失，還包含蒸餾損失。這里的知識蒸餾則與傳統(tǒng)知識蒸餾類似，是為了將解碼器的特征知識遷移至小網(wǎng)絡(luò)。

在蒸餾過程中計算中間層特征的分布，將特征圖的分布信息定義為知識，通過最大均值差異(max mean discrepancy，MMD)[12]衡量大網(wǎng)絡(luò)(解碼器)和小網(wǎng)絡(luò)之間的特征分布差異。以MMD作為蒸餾損失函數(shù)，引導(dǎo)小網(wǎng)絡(luò)中間層的激活分布模擬大網(wǎng)絡(luò)的激活分布。

(4)

(5)

(5)式中：‖·‖2是L2正則化；G是Gram矩陣，矩陣中每一項為gij=(fi)Tfj。Gram矩陣是特征圖向量化后內(nèi)積的結(jié)果，能反應(yīng)特征之間的相關(guān)程度。

小網(wǎng)絡(luò)總的損失函數(shù)LS可表示為

LS=Lsr+λLF+βSLdistill

(6)

(6)式中：Lsr為超分辨率任務(wù)常用的重建損失，由生成圖像與真實高分辨率圖像計算MAE損失而得；λ是退化損失函數(shù)的平衡系數(shù)，與大網(wǎng)絡(luò)中的參數(shù)設(shè)置相同，λLF為退化損失，與解碼器相同，將生成圖像經(jīng)退化后與輸入圖像計算MAE損失；βSLdistill為特征蒸餾損失，其中βS為蒸餾損失的平衡系數(shù)。

3 實 驗

3.1 數(shù)據(jù)集和評價指標(biāo)

參照EDSR的訓(xùn)練設(shè)置[14-15]，本文使用DIV2K數(shù)據(jù)集進行訓(xùn)練，其中包含800張高分辨率圖像，低分辨率圖像通過對高分辨圖像進行雙三次下采樣生成而得。每個訓(xùn)練批次大小為32張圖片。

本文在Set5、Set14、B100和Urban100等標(biāo)準(zhǔn)數(shù)據(jù)集上評估提出的方法，并使用亮度通道上計算的PSNR作為評估指標(biāo)。

3.2 實驗分析

首先訓(xùn)練大網(wǎng)絡(luò)，經(jīng)過反復(fù)實驗，設(shè)置超參數(shù)βT=10-4,βS=10-3,λ=0.1。表2所示為在不同βT值情況下訓(xùn)練小網(wǎng)絡(luò)，并在Set5數(shù)據(jù)集上驗證兩倍超分的結(jié)果。在這部分實驗中，退化支路沒有引入其中，以免對編解碼器的調(diào)參實驗造成影響。當(dāng)βT為0時，編解碼器失去了壓縮HR中高頻特征的功能，成為了簡單的線性映射。若參數(shù)βT設(shè)置太大，編碼器損失函數(shù)將促使編碼器生成的低維特征與低分辨率圖像趨于同質(zhì)。在這種情況下，本文提出的知識蒸餾框架將不能從高分辨率圖像的先驗知識中獲益，并且小網(wǎng)絡(luò)的特征蒸餾與傳統(tǒng)蒸餾方法毫無差別，性能提升微乎其微。

表2 超參數(shù)βT的實驗結(jié)果Tab.2 Results of balance parameters βT

經(jīng)過大量實驗得知，參數(shù)βT設(shè)置為10-4時，能在保證解碼器學(xué)習(xí)超分辨率映射和編碼器提取高分辨率的先驗知識之間取得折衷，能通過特征蒸餾使小網(wǎng)絡(luò)達(dá)到最佳效果。

小網(wǎng)絡(luò)在不同數(shù)據(jù)集上進行2倍放大的蒸餾結(jié)果如表3所示，進行4倍放大結(jié)果如表4所示。經(jīng)過蒸餾訓(xùn)練后，小網(wǎng)絡(luò)在各個基準(zhǔn)數(shù)據(jù)集上的性能都有所提升，在網(wǎng)絡(luò)參數(shù)量和計算復(fù)雜度未增加的情況下，依靠本文提出的知識蒸餾方法，小網(wǎng)絡(luò)PSNR在進行2倍超分辨率時能提升0.17～0.28 dB；在4倍超分辨率時能提升0.11～0.18 dB。與傳統(tǒng)知識蒸餾方法相比，本文方法大大提升了小網(wǎng)絡(luò)性能，能將知識蒸餾高效地運用在超分辨率任務(wù)上。這主要有以下兩個原因：①大網(wǎng)絡(luò)中編解碼器結(jié)構(gòu)能有效捕捉高分辨率圖像中的高頻信息，并通過特征蒸餾傳遞給小網(wǎng)絡(luò)；②退化支路的約束縮小了超分辨率任務(wù)的解空間，加速了小網(wǎng)絡(luò)的收斂。傳統(tǒng)知識蒸餾方法未對大網(wǎng)絡(luò)進行專門化訓(xùn)練，使得其蘊含的高頻信息包含太多噪聲，無法指導(dǎo)小網(wǎng)絡(luò)訓(xùn)練。

表3—表4中的SRKD[18]和PISR[19]也是兩種結(jié)合了知識蒸餾的超分辨率方法。其中，SRKD是基于傳統(tǒng)知識蒸餾AT[9]的方式，通過學(xué)習(xí)大網(wǎng)絡(luò)的特征圖進行蒸餾，只是修改了AT中的知識定義，因此，它并不能脫離傳統(tǒng)知識蒸餾在超分任務(wù)上的桎梏，PISR類似于本文方法，通過對小網(wǎng)絡(luò)的預(yù)訓(xùn)練提取高分辨率圖像中原有的先驗信息來提高蒸餾效果。本文按照文獻[18-19]的參數(shù)設(shè)置，在DIV2K數(shù)據(jù)集上完成了復(fù)現(xiàn)。結(jié)果表明，無論是在2倍還是在4倍超分辨率上，本文方法都取得了更高的PSNR指標(biāo)。

表3 小網(wǎng)絡(luò)在不同數(shù)據(jù)集上進行2倍放大的蒸餾結(jié)果Tab.3 Small network distillation results with 2x magnification on different data sets

表4 小網(wǎng)絡(luò)在不同數(shù)據(jù)集上進行4倍放大的蒸餾結(jié)果Tab.4 Small network distillation results with 4x magnification on different data sets

將本文方法訓(xùn)練而得的小網(wǎng)絡(luò)和文獻[20-21]的超分辨率網(wǎng)絡(luò)進行對比，從表3—表4可得，DRCN和MemNet的參數(shù)量都大于小網(wǎng)絡(luò)，相應(yīng)地計算復(fù)雜度更高，因而在各個數(shù)據(jù)集上它們都取得了比未蒸餾小網(wǎng)絡(luò)更高的PSNR結(jié)果。但經(jīng)過本文方法的訓(xùn)練，在不需對網(wǎng)絡(luò)結(jié)構(gòu)進行特殊設(shè)計情況下，蒸餾后的小網(wǎng)絡(luò)并未增加運算復(fù)雜度或提高網(wǎng)絡(luò)參數(shù)量，它的性能表現(xiàn)已經(jīng)超過了DRCN，同時在耗時上也縮減至與MemNet相近的水準(zhǔn)，這說明本文的知識蒸餾方法能針對超分辨率網(wǎng)絡(luò)實現(xiàn)良好的輕量化效果。

3.3 消融實驗

為了充分證明有效性，本文還對方法中的編解碼器結(jié)構(gòu)、退化支路、蒸餾損失函數(shù)進行了消融實驗。在Set5上進行2倍放大，對PSNR進行定量計算分析，比較其超分辨率性能。

表5所示為在不同模塊組合下訓(xùn)練而得小網(wǎng)絡(luò)的消融實驗結(jié)果。

表5 消融實驗的結(jié)果Tab.5 Results of ablation studies

從表5可以看出，相比于傳統(tǒng)蒸餾方法，編解碼器結(jié)構(gòu)使得特征蒸餾在超分辨率任務(wù)上更加有效，小網(wǎng)絡(luò)受益于大網(wǎng)絡(luò)解碼器的網(wǎng)絡(luò)權(quán)重，這為小網(wǎng)絡(luò)提供了一個良好的訓(xùn)練起點，并且遷移了大網(wǎng)絡(luò)的重構(gòu)能力；MMD損失函數(shù)提供了比MAE更好的結(jié)果，基于MMD的蒸餾損失促使小網(wǎng)絡(luò)和大網(wǎng)絡(luò)的特征圖保持一致分布，一定程度上避免了大網(wǎng)絡(luò)特征圖中的噪聲影響；退化支路提供了更好的性能表現(xiàn)，能帶來0.1 dB的PSNR提升，這表明縮小解空間有助于超分辨率任務(wù)。

4 結(jié) 論

本文主要研究了如何使用知識蒸餾對超分辨率網(wǎng)絡(luò)進行輕量化，全面分析了現(xiàn)有知識蒸餾方法直接運用到超分辨率網(wǎng)絡(luò)的局限性，并提出了一種基于編解碼器的知識蒸餾訓(xùn)練框架，能有效提取高分辨率圖像中的先驗知識，再通過特征蒸餾將其傳遞給小網(wǎng)絡(luò)。本文方法顯著地提高了超分辨率網(wǎng)絡(luò)的性能，實現(xiàn)了超分辨率網(wǎng)絡(luò)的輕量化。