高效單圖像超分辨率重建:深監(jiān)督對稱蒸餾網(wǎng)絡(luò)

毛盼娣,徐道連

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

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

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

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

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

1 相關(guān)工作

1.1 單圖像超分辨率

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

1.2 知識蒸餾

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

本文方法相關(guān)方法包括SRKD,FAKD和PISR.SRKD對教師網(wǎng)絡(luò)和學(xué)生網(wǎng)絡(luò)不同階段進(jìn)行蒸餾,考慮到教師網(wǎng)絡(luò)與學(xué)生網(wǎng)絡(luò)通道數(shù)不同,SRKD對中間特征的統(tǒng)計信息進(jìn)行監(jiān)督.FAKD改進(jìn)了SRKD,在整個蒸餾的過程中FAKD與SRKD類似,區(qū)別在于FAKD在蒸餾損失方面提出了特征關(guān)聯(lián)機(jī)制,取得了不錯的效果.最近PISR被提出,PISR首次將HR作為特權(quán)信息引入SISR.PISR與SRKD、FAKD有本質(zhì)的不同,其將HR直接輸入教師網(wǎng)絡(luò),允許教師網(wǎng)絡(luò)提取更多的信息轉(zhuǎn)移到學(xué)生網(wǎng)絡(luò).

1.3 深監(jiān)督

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

2 方 法

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

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

2.1 教師網(wǎng)絡(luò)

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

2.1.1高頻特征遞歸模塊

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

(1)

(2)

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

2.1.2重建網(wǎng)絡(luò)

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

(3)

2.1.3對稱退化模塊

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

(4)

(5)

2.2 深監(jiān)督

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

(6)

(7)

(8)

(9)

(10)

(11)

(12)

(13)

2.3 學(xué)生網(wǎng)絡(luò)

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

(14)

初始化網(wǎng)絡(luò)權(quán)重對學(xué)生網(wǎng)絡(luò)的重建性能至關(guān)重要[14].使用教師網(wǎng)絡(luò)的蒸餾知識初始化學(xué)生網(wǎng)絡(luò),可以將教師網(wǎng)絡(luò)訓(xùn)練好的重建能力轉(zhuǎn)移給學(xué)生網(wǎng)絡(luò),為后期優(yōu)化提供了很好的起點(diǎn).

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

(15)

(16)

蒸餾損失將教師網(wǎng)絡(luò)的蒸餾知識轉(zhuǎn)移到學(xué)生網(wǎng)絡(luò),本文同PISR使用VID.目的是最大化教師網(wǎng)絡(luò)和學(xué)生網(wǎng)絡(luò)對應(yīng)特征信息熵之間的互信息,

(17)

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

(18)

3 實(shí) 驗

3.1 實(shí)驗細(xì)節(jié)

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

3.2 實(shí)驗結(jié)果

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

3.3 消融實(shí)驗

消融實(shí)驗對DSSD每個組件進(jìn)行對比實(shí)驗,以驗證每個模塊的作用,組合在一起之后會達(dá)到最好的效果.

3.3.1高頻遞歸模塊實(shí)驗

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

表1 HFRM模塊的結(jié)果比較

3.3.2對稱退化模塊實(shí)驗

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

表2 SR的退化方式對比實(shí)驗

3.3.3深監(jiān)督實(shí)驗

表3 監(jiān)督對SR性能的影響

3.3.4教師網(wǎng)絡(luò)通用性對比實(shí)驗

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

表4 其他SR方法的定量結(jié)果

3.3.5視覺效果

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

4 結(jié) 論

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

附 錄

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