高效單圖像超分辨率重建:深監督對稱蒸餾網絡

毛盼娣,徐道連

(1.重慶城市科技學院 電氣工程與智能制造學院,重慶 402167;2.重慶大學 光電工程學院,重慶 400044)

單圖像超分辨率(SISR)旨在利用低分辨率(Low-Resolution,LR)圖像重建出對應的高分辨率(HR)圖像,近年來該方法在物體檢測[1]、信息取證[2]等領域應用十分廣泛.DONG等[3]提出SRCNN首次為SISR引入了卷積神經網絡(CNNs).隨后,基于CNNs的SISR方法致力于設計更深更廣的網絡架構以實現性能提高.然而這類方法往往參數量和計算量巨大,例如EDSR[4]和RCAN[5]等.在現實場景下真正需求的是輕量化網絡模型,尤其是可以部署到移動設備的實時模型,因而如何更好地實現網絡輕量化是一個亟待解決的問題.

SISR模型輕量化首先想到的便是設計特殊的模型,減少網絡的層數、使用小的卷積核、采用輕量化組件,例如分組卷積[6]和可變形卷積[7]等.但是特殊設計的模型依舊存在問題:一是輕量級網絡與EDSR和RCAN等大網絡相比具有很大的差距;二是輕量級模型中一些組件在實際硬件芯片上不能被很好地支持.SISR領域的蒸餾方法應運而生,其使用性能好的大模型(教師網絡)蒸餾信息提升小模型(學生網絡),使學生網絡也能獲得接近教師網絡的性能.本文基于第二種方法探索.

SISR領域的蒸餾方法中,首次提出的SRKD[6]對教師網絡和學生網絡的不同階段進行蒸餾,考慮到教師網絡和學生網絡通道數不相同,SRKD對中間信息的統計進行監督.隨后,FAKD在SRKD的基礎上引入了特征關聯機制[8],進一步提升了學生網絡的性能.但這些方法仍有一定的局限性,于是PISR首次引入特權信息(HR圖像)[9],用教師網絡提取HR高頻特征補充到學生網絡,明顯提升了重建性能.利用特權信息的方法最近得到了明顯的關注,基于參考圖像的超分辨率方法C2-Matching[10]利用特權信息大幅提升了性能(目前最好的RefSR模型).明顯地,教師網絡的性能決定了學生網絡的上限,但是現有利用特權信息的方法存在很大的問題.

根據局部相似性,教師網絡提取HR高頻特征因沒有很好的約束導致學生網絡對輸入互補性不夠準確,如圖1所示.針對這個問題,本文提出了高效單圖像超分辨率重建[11]:深監督對稱蒸餾網絡(DSSD).DSSD針對教師網絡提取HR的高頻特征不夠準確提出了一個新穎的教師網絡.為了約束中間層構造了高頻特征遞歸模塊(HFRM),這種遞歸方式對中間特征約束更強.為了使教師網絡學習到足夠準確的高頻特征,提出了一種對稱退化結構(SDM),SDM與HFRM一樣進一步約束了教師網絡中的SR.特別地,整個教師網絡使用深監督加強對教師網絡的約束,同時解決了教師網絡難以訓練的問題,使轉移到學生網絡的蒸餾知識更加準確.同時,本文為了更好適用于移動設備,選擇只有13 k參數量和6 GFlops計算量的FSRCNN[5]網絡.本文主要貢獻:1)提出了一種新穎的通用蒸餾教師網絡架構DSSD,該架構可以更準確地提取HR高頻細節.2)其中構造了HFRM和SDM模塊,目的是使教師網絡更準確地提取HR高頻特征.這兩個模塊采取遞歸方式并且結構簡單.3)為了約束教師網絡中提取的高頻特征,采用深監督方法使教師網絡蒸餾的知識可以與學生網絡互補.

1 相關工作

1.1 單圖像超分辨率

隨著CNNs的發展,DONG等[3]首次提出基于CNNs的SISR方法SRCNN學習LR到HR的映射關系.之后的方法使用大量卷積層,利用密集和跳躍連接等保留更多的特征[12],與之前方法相比有了顯著的提升.為了減少顯存和運行時間,高效SISR被提出.對于顯存高效方法,其本質是利用特殊設計的架構減少參數量.對于高效運行時間方法,其本質是較低的計算量,使用級聯、多分支架構或組卷積等方式.但是這些特殊設計的高效SISR方法不能適用于移動設備.

1.2 知識蒸餾

知識蒸餾的提出是為了轉移大模型(教師網絡)軟標簽到小模型(學生網絡)以增強小模型性能,它被廣泛運用在網絡壓縮上[13].之后特征蒸餾被提出,可以從教師網絡轉移特征至學生網絡,幫助學生網絡更好地訓練.利用特權信息的方法本質是使用額外信息,這在訓練時需要額外成本但是在測試時不需要,被稱之為廣義蒸餾.廣義蒸餾能夠從教師網絡轉移到學生網絡更多的信息以增強學生網絡.特別地,變分信息蒸餾VID通過最大化教師和學生網絡特征之間的互信息來轉移知識,本文同PISR一樣使用VID.

本文方法相關方法包括SRKD,FAKD和PISR.SRKD對教師網絡和學生網絡不同階段進行蒸餾,考慮到教師網絡與學生網絡通道數不同,SRKD對中間特征的統計信息進行監督.FAKD改進了SRKD,在整個蒸餾的過程中FAKD與SRKD類似,區別在于FAKD在蒸餾損失方面提出了特征關聯機制,取得了不錯的效果.最近PISR被提出,PISR首次將HR作為特權信息引入SISR.PISR與SRKD、FAKD有本質的不同,其將HR直接輸入教師網絡,允許教師網絡提取更多的信息轉移到學生網絡.

1.3 深監督

深監督是在訓練神經網絡時,針對中間隱藏層特征透明度不高以及深層網絡難以訓練的問題,對隱藏層進行再監督,促進網絡更快更好收斂.DSN首次提出深監督概念,但DSN缺點在于使用的網絡結構不夠深,為了更好地訓練深度網絡,通過給神經網絡的某些層添加了輔助監督來解決這個問題.此外,應用于分割的Unet++也適用深監督增強醫學圖像分割的性能.由于教師網絡中間層較多,為了更好約束中間層,使中間特征能更準確地蒸餾到學生網絡,本文采用深監督的方式,明顯提升了DSSD的性能.

2 方 法

本文提出了一個新穎的深監督對稱蒸餾教師網絡.該網絡利用特權信息訓練教師網絡,訓練好的教師網絡會幫助學生網絡更好地重建.由于篇幅有限,本文以倍率4為例:使用X和Y表示LR和HR圖像.其中′表示二分之一HR大小,″表示四分之一HR大小,ζ表示監督函數.

這種利用特權信息的對稱結構類似于在課堂學習(訓練階段)時,不僅僅會有正確答案(HR),而且老師會解釋其中的來龍去脈(HFRM),更進一步老師會告知錯在哪里(SDM)以及該怎么修正(深監督).

2.1 教師網絡

學生網絡性能和教師網絡性能直接相關,所以提高教師網絡性能是關鍵.為此本文設計了一種對稱教師網絡結構如圖2所示,該網絡輸入與PISR相同(HR).使用HFRM模塊逐級下采樣,每經過一個HFRM模塊分辨率縮小二分之一,在遞歸2個HFRM模塊后變成了HR分辨率的四分之一,與LR尺寸相同可以輸入FSRCNN.在經過FSRCNN之后生成SR,將SR也采用與HFRM相同的下采樣可以保證SR蘊含更多的高頻特征,于是構造了SDM模塊.特別地,SDM與HFRM模塊在結構上相同,形成一種對稱網絡結構.同時為了轉移更準確的高頻特征,教師網絡采取深監督的方式對中間層進行約束,取得了不錯的效果.

2.1.1高頻特征遞歸模塊

本文構造了HFRM模塊將HR遞歸投影到低維特征空間生成緊湊的高頻特征.與機械式雙三次下采樣不同,基于CNNs的HFRM能學習到更多與LR互補的高頻特征,然后從緊湊特征重構SR圖像.遞歸HFRM受到深監督的約束可以提取更準確的高頻特征.這一過程表示為式(1)、(2):

(1)

(2)

HFRM模塊詳細的結構如圖3所示,k3n56s1為先經過一個3×3的卷積將3通道擴展到56(為了與FSRCNN保持一致),步長大小為1,激活函數為PReLU.第2個3×3的卷積步長為2進行下采樣,之后再有兩個3×3的卷積,先將56通道降至12通道,再變為3通道.經過2個HFRM生成與LR大小一致的緊湊特征交給FSRCNN進行重構,其中遞歸能很好地受到深監督約束.

2.1.2重建網絡

本文目的是構造適用于移動設備的網絡,既要很小的參數量也要很小的計算量,同時避免特殊設計的網絡架構,所以FSRCNN是一個好的選擇.教師網絡中重建網絡過程如式(3)所示:

(3)

2.1.3對稱退化模塊

為了保證經過FSRCNN重建的SR保留足夠準確的高頻特征,本文構造了SDM模塊(圖4),之前大部分方法主要是最小化SR與HR之間的損失而忽略了解空間的問題.對于LR來說,LR與SR是一種一對多的關系,利用SDM模塊去約束其解空間的大小,保證SR具有更準確的高頻特征.特別為了保持簡單的結構,SDM模塊與HFRM模塊有著相同的構造,同樣采用遞歸下采樣的方式并受到深監督的約束,如式(4)、(5):

(4)

(5)

2.2 深監督

DSSD除了新穎的對稱架構以外,還有一個明顯的特點:它由多個相同的監督構成深監督約束.如圖2所示,每個HFRM模塊可以下采樣2倍,于是4倍率需要遞歸2次.HR圖像在經過HFRM模塊之后會保留更多高頻特征,這種特征會轉移到學生網絡.但是在保留的過程中,注意到學生網絡的輸入是雙三次下采樣得到的LR圖像,盡管盡可能多地保留高頻特征對學生網絡有用,但是總體上不宜與LR偏差太大(局部相似性).為了解決這一個問題,本文使用深監督策略,每一級分辨率都受到監督,保證在保留高頻特征的同時保證合理性.

(6)

(7)

(8)

(9)

(10)

(11)

(12)

(13)

2.3 學生網絡

學生網絡擁有與教師網絡中重建網絡(FSRCNN)相同的結構,但是輸入不同.學生網絡的輸入為雙三次下采樣的LR,然后重建SR,如式(14):

(14)

初始化網絡權重對學生網絡的重建性能至關重要[14].使用教師網絡的蒸餾知識初始化學生網絡,可以將教師網絡訓練好的重建能力轉移給學生網絡,為后期優化提供了很好的起點.

盡管教師網絡的蒸餾知識可以很好地初始化學生網絡,但是由于學生網絡的輸入不同,因此需要使用重建損失和蒸餾損失進一步優化學生網絡.其中式(16)中λS為10-6.

(15)

(16)

蒸餾損失將教師網絡的蒸餾知識轉移到學生網絡,本文同PISR使用VID.目的是最大化教師網絡和學生網絡對應特征信息熵之間的互信息,

(17)

ζ(x)=ln (1+ex).

(18)

3 實 驗

3.1 實驗細節

本文使用DIV2K數據集訓練,該數據集由800對LR和HR圖像組成,其中LR是由HR雙三次下采樣得到.隨機將HR圖像剪切為192×192大小的塊,同樣的LR塊取自對應比例因子的LR圖像.同時使用數據增強技術,包括隨機旋轉和水平翻轉.教師網絡使用隨機初始化訓練,訓練模型時,批處理大小為16,一共訓練1 000輪,其中β1=0.900,β2=0.999.對于學習率大小,設置為10-3并且使用余弦退火技術將其減少到10-5.使用標準評估集,用峰值信號比(PSNR)和結構相似程度(SSIM)在Y通道上評估.本文運行于NVIDIA Titan RTX.

3.2 實驗結果

附表Ⅰ比較了DSSD學生模型與最先進模型的性能,特別是高效SISR方法(*表示使用DIV2K重新訓練,最好和最壞結果分別用下劃實線和下劃曲線標出).為了進行定量對比,評價指標為標準數據集在2、3、4×下的PSNR和SSIM,也包含模型的參數量計算量運行時間,其中計算量是在1 280×720大小的HR圖像上測量.從附表Ⅰ可以看到:DSSD學生模型在所有比例因子上均優于PISR,甚至4倍率Set5超PISR接近0.12 dB,但是大數據集提升不如小數據集.因FSRCNN只有13 k參數量和6 GFlpos計算量,其重建能力在大數據集上已經捉襟見肘.同理,DSSD在較小的比例因子下的提升也有限,是因為較小的比例因子下丟失的信息較少,小參數量的模型重建能力有限.針對這個問題,本文后續在3.3.4討論了較大模型的性能提升.

3.3 消融實驗

消融實驗對DSSD每個組件進行對比實驗,以驗證每個模塊的作用,組合在一起之后會達到最好的效果.

3.3.1高頻遞歸模塊實驗

為了驗證HFRM的有效性,設計了表1對比實驗:這里只探究HFRM模塊,在設計此對比實驗時去掉SDM模塊以保證公平性.可以看到直接雙三次下采樣效果最差,因為機械式下采樣就等于直接輸入LR圖像,導致教師網絡和學生網絡的輸入相同.FSRCNN的PSNR為30.89 dB,直接雙三次下采樣方法訓練本質上是用教師網絡訓練好的特征初始化學生網絡.PISR的Encoder使用CNNs學到一部分LR沒有的高頻特征,但轉移的高頻特征不夠準確(如圖5對比所示).HFRM模塊采用遞歸的方式約束中間特征,目的是可以學習到更準確更合理的高頻特征.

表1 HFRM模塊的結果比較

3.3.2對稱退化模塊實驗

為了驗證SDM的有效性,本文設計了表2對比實驗.這里保留HFRM只改變SDM模塊以獲得公平的對比環境.從表2可以看出,去掉SDM就是表1最后一個實驗.重點在表2的后3個消融實驗:將SDM使用雙三次下采樣結果反而變差,原因是這種機械方式破壞了SR的退化結果,深監督會使SR向更壞的方向發展.而基于CNNs下采樣限制SR的高頻特征提高了SR的重建性能.最后是本文構造的SDM模塊,在深監督的配合下進一步約束了SR高頻信息的合理性,獲得了更好的重建效果.

表2 SR的退化方式對比實驗

3.3.3深監督實驗

表3 監督對SR性能的影響

3.3.4教師網絡通用性對比實驗

為了探究DSSD教師網絡是否具有通用性,本文挑選了幾個參數量較大的輕量級SISR方法,見表4.圖6為標準測試集在4視覺上的比較為665 k和591 k,*表示用DIV2K訓練不設置蒸餾,可以看見DSSD在SISR方法都有提升.特別地,作用在較小的數據集Set5上提升較大,而在較大數據集上B100提升有限,這是因為FSRCNN參數量只有13 k已經到了性能極限.在增大參數量的網絡VDSR上DSSD增強了更多,接近0.18 dB(Set5),而在IDN上的大數據集B100上增加了0.1 dB.這些實驗驗證了DSSD教師網絡具有通用性.

表4 其他SR方法的定量結果

3.3.5視覺效果

圖6展示了使用學生網絡在標準測試集上的重建結果.可以清楚地看到學生模型提供了比原始基線模型(FSRCNN)和目前SOTA蒸餾方法PISR更好的性能.提出DSSD的目的在于探究如何更好地提取HR中包含的高頻紋理細節,這些高頻紋理細節可以使圖像的邊緣更清晰更明顯,可以明顯看見B100上數據集少了PISR的棋盤偽影,同時Urban100上數據集邊緣明顯變清晰銳利,Set14中barbara的重建效果明顯紋理細節更多.這些效果圖證明了DSSD確實具有很好提取HR高頻細節的能力.

4 結 論

本文提出了一種新穎的高效單圖像超分辨率重建方法:深監督對稱蒸餾網絡.針對教師網絡提取HR高頻信息不夠準確這一問題,構造了高頻特征遞歸模塊和對稱退化模塊,有效提高了重建性能.在整個蒸餾的教師網絡中,利用深監督能更好地約束中間特征使提取的特征更準確.消融實驗很好地證明了本文方法的有效性,將在未來的工作中繼續探索蒸餾方法中特征轉移部分以進一步提高蒸餾方法的性能.

附 錄

附表Ⅰ見電子版(DOI:10.16366/j.cnki.1000-2367.2023.06.007).