基于CNN-Transformer混合構(gòu)架的輕量圖像超分辨率方法*

林承浩，吳麗君

(福州大學(xué) 物理與信息工程學(xué)院，福建 福州 350108)

0 引言

圖像超分辨率(Super Resolution，SR)是一項(xiàng)被廣泛關(guān)注的計(jì)算機(jī)視覺任務(wù)，其目的是從低分辨率(Low Resolution，LR)圖像中重建出高質(zhì)量的高分辨率(High Resolution，HR)圖像[1]。由于建出高質(zhì)量的高分辨率圖像具有不適定的性質(zhì)，因此極具挑戰(zhàn)性[2]。隨著深度學(xué)習(xí)等新興技術(shù)的崛起，許多基于卷積神經(jīng)網(wǎng)絡(luò)(CNN)的方法被引入到圖像超分任務(wù)中[3-6]。SRCNN[3]首次將卷積神經(jīng)網(wǎng)絡(luò)引入到圖像超分任務(wù)中，用卷積神經(jīng)網(wǎng)絡(luò)來學(xué)習(xí)圖像的特征表示，并通過卷積層的堆疊來逐步提取更高級別的特征，使得重建出的圖像具有較高的質(zhì)量。在后續(xù)研究中，Kaiming He等人提出了殘差結(jié)構(gòu)ResNet[5]，通過引入跳躍連接，允許梯度能夠跨越層進(jìn)行傳播，有助于減輕梯度消失的問題，使得模型在較深的網(wǎng)絡(luò)情況下仍然能保持較好的性能。Bee Lim等人在EDSR[6]中也引入了殘差結(jié)構(gòu)，EDSR實(shí)際上是SRResnet[7]的改進(jìn)版，去除了傳統(tǒng)殘差網(wǎng)絡(luò)中的BN層，在節(jié)省下來的空間中擴(kuò)展模型尺寸來增強(qiáng)表現(xiàn)力。RCAN[8]中提出了一種基于Residual in Residual結(jié)構(gòu)(RIR)和通道注意力機(jī)制(CA)的深度殘差網(wǎng)絡(luò)。雖然這些模型在當(dāng)時(shí)取得了較好的效果，但本質(zhì)上都是基于CNN網(wǎng)絡(luò)的模型，網(wǎng)絡(luò)中卷積核的大小會限制可以檢測的空間范圍，導(dǎo)致無法捕捉到長距離的依賴關(guān)系，意味著它們只能提取到局部特征，無法獲取全局的信息，不利于紋理細(xì)節(jié)的恢復(fù)，使得圖像重建的效果不佳[5]。

由于Transformer在自然語言處理(Natural Language Processing，NLP)領(lǐng)域中取得了較好的成效[9]，Alexey Dosovitskiy等人將Transformer引入到計(jì)算機(jī)視覺(CV)領(lǐng)域中，即ViT(Vision Transformer)，并且在多個(gè)視覺任務(wù)中取得了成功[10]。ViT的優(yōu)勢在于Transformer對于全局的信息更加敏感，模型中的注意力模塊能夠在輸入序列的所有位置上進(jìn)行全局交互，從而捕捉到長距離的依賴關(guān)系。但是，用于圖像超分辨率的ViT，需要將輸入圖像分割成固定大小的塊，并對每個(gè)小塊進(jìn)行獨(dú)立的處理[11]，這種處理策略就會產(chǎn)生兩個(gè)弊端：一是修復(fù)后的圖像可能會在每個(gè)小塊周圍引入邊界偽影；二是每個(gè)補(bǔ)丁的邊界像素可能會丟失信息，影響到重建圖像的質(zhì)量。Ze Liu等人提出了Swin-Transformer[11]，將滑動窗口機(jī)制引入到Transformer中，其中滑窗操作包括了不重疊的Local-window和重疊的Cross-window，將注意力計(jì)算限制在一個(gè)窗口內(nèi)可以大幅度節(jié)省計(jì)算量，并且通過滑窗操作也能使得注意力機(jī)制能夠注意到全局的特征。

Jingyun Liang等人結(jié)合卷積神經(jīng)網(wǎng)絡(luò)和Swin-Transformer兩者的優(yōu)點(diǎn)提出了SwinIR[12]，將兩個(gè)構(gòu)架以串行的方式應(yīng)用在超分領(lǐng)域，展示出了混合構(gòu)架在超分任務(wù)中的巨大應(yīng)用前景。Peng等人提出了一種以CNN-Transformer并行的方法Conformer[13]，它由一個(gè)CNN分支和一個(gè)Transformer分支組成，依靠特征耦合單元(Feature Coupling Unit，F(xiàn)CU)以交互的方式在不同分辨率下融合局部特征和全局特征，其結(jié)果表明了并行結(jié)構(gòu)能以最大限度地保留局部特征和全局特征。雖然Conformer結(jié)合了兩個(gè)網(wǎng)絡(luò)的優(yōu)勢，但訓(xùn)練的參數(shù)量和訓(xùn)練時(shí)長也相應(yīng)地增加了。并且，圖像超分任務(wù)通常需要輸入較高分辨率的圖像，占用大量GPU內(nèi)存，限制了模型的靈活性。要想取得更加高清的圖像勢必會增加網(wǎng)絡(luò)模型中的參數(shù)量[14]。因此，在結(jié)合兩種構(gòu)架優(yōu)勢的同時(shí)，降低訓(xùn)練成本，使得圖像超分辨率模型輕量化成為了本文需要解決的重大問題。

針對上述問題，本文提出了基于Swin-Transformer的單圖像超分網(wǎng)絡(luò)STSR(Swin-Transformer-based Single Image Super-Resolution)。具體貢獻(xiàn)如下：結(jié)合CNN和Transformer的優(yōu)勢，設(shè)計(jì)了并行特征提取的特征增強(qiáng)模塊(Feature Enhancement Block，F(xiàn)EB)，能夠有效地捕捉圖像的局部細(xì)節(jié)特征，同時(shí)也能夠捕捉長距離的依賴關(guān)系，使模型具有全局上下文建模的能力。本文采用了輕量化的Transformer模塊，在達(dá)到較好重建效果的同時(shí)還能保持較低的計(jì)算成本。此外，通過設(shè)計(jì)動態(tài)調(diào)整模塊(Dynamic Adjustment，DA)，可根據(jù)輸入圖像的特征對輸出進(jìn)行動態(tài)的調(diào)整，從而增強(qiáng)網(wǎng)絡(luò)的擬合能力，使得重建圖像的紋理細(xì)節(jié)更加貼近于真實(shí)圖像。

1 本文算法

本文設(shè)計(jì)的基于Swin-Transformer的單圖像超分網(wǎng)絡(luò)STSR的整體網(wǎng)絡(luò)結(jié)構(gòu)如圖1所示。結(jié)合CNN和Transformer兩者的優(yōu)勢，本文設(shè)計(jì)了一個(gè)并行特征提取的特征增強(qiáng)模塊FEB，采用卷積神經(jīng)網(wǎng)絡(luò)和Transformer網(wǎng)絡(luò)并行提取圖像特征，為了減少訓(xùn)練中的參數(shù)量，F(xiàn)EB中的輕量化Transformer模塊只采用Swin-Transformer中的STL(Swin Transformer Layer)來提取特征，并且在網(wǎng)絡(luò)中堆疊3層FEB以獲取更深層次的特征，有助于提升重建圖像的質(zhì)量。為了進(jìn)一步的提升重建圖像的質(zhì)量，本文設(shè)計(jì)了一個(gè)動態(tài)調(diào)整模塊DA，動態(tài)調(diào)整模塊主要采用的是通道注意力機(jī)制，通過這一機(jī)制，網(wǎng)絡(luò)可以根據(jù)輸入圖像的內(nèi)容和結(jié)構(gòu)，自適應(yīng)地增強(qiáng)或者減弱不同通道的影響，也可以減少網(wǎng)絡(luò)對無關(guān)信息的依賴，使得在訓(xùn)練過程中網(wǎng)絡(luò)能夠快速收斂，減少訓(xùn)練時(shí)間，同時(shí)能夠使得重建圖像更加貼近于真實(shí)圖像。STSR主要包括了四個(gè)模塊：淺層特征提取、深層特征提取，動態(tài)調(diào)整模塊和高質(zhì)量圖像重建模塊。

圖1 STSR的整體網(wǎng)絡(luò)框架

(1)淺層特征提取

首先，使用一個(gè)3×3的卷積層來提取圖像的淺層特征F0，如式(1)所示。

F0=H3×1(ILR)

(1)

(2)深層特征提取

因?yàn)闇\層特征提取模塊只經(jīng)過一次簡單的卷積操作，提取到較為初級的特征。將提取到的淺層特征F0，使用深層特征提取模塊進(jìn)一步地提取圖像特征。深層特征提取模塊由K個(gè)FEB模塊和一個(gè)3×3的卷積層構(gòu)成，用來提取深層特征FDF，如式(2)所示。

FDF=HDF(F0)

(2)

其中每個(gè)FEB的輸出F1、F2、FK以及輸出的深層特征FDF，如式(3)(4)所示。式中HKFi表示第i個(gè)FEB模塊，Hconv表示最終的卷積層。

Fi=HKFi(Fi-1)，i=1，2，…，K

(3)

FDF=Hconv(FK)

(4)

(3)動態(tài)調(diào)整模塊

動態(tài)調(diào)整模塊主要采用通道注意力機(jī)制，通過計(jì)算通道間的相關(guān)性和重要性，網(wǎng)絡(luò)可以自適應(yīng)地增強(qiáng)或減弱不同通道的影響。這有助于提取輸入圖像中最相關(guān)和有用的特征，使得重建圖像與輸入圖像具有更高的相似度，從而改善超分辨率重建的效果。

動態(tài)調(diào)整模塊接收來自淺層特征提取模塊輸出的淺層特征F0，經(jīng)過處理后得到動態(tài)調(diào)整特征FDA，如式(5)所示。

FDA=HDA(F0)

(5)

(4)高質(zhì)量圖像重建模塊

圖像重建模塊其實(shí)就是卷積與上采樣的組合，本文中采用的是“卷積+pixel shuffle+卷積”的方式來進(jìn)行圖像重建。作為網(wǎng)絡(luò)的最后部分，將接收淺層特征F0，深層特征FDF和動態(tài)調(diào)整特征FDA，以獲得重建圖像ISR。f和fp分別代表的是卷積層和亞像素卷積層，計(jì)算ISR的公式如下：

ISR=f(fp(f(FDF)))+f(fp(f(FDA)))+f(fp(F0))

(6)

1.1 特征增強(qiáng)模塊(FEB)

本文方法中提出的FEB模塊主要是由卷積神經(jīng)網(wǎng)絡(luò)CNN、輕量型Transformer和特征融合模塊三個(gè)部分組成。特征增強(qiáng)模塊FEB的大致結(jié)構(gòu)如圖2(a)所示。

圖2 FEB模塊和STL模塊結(jié)構(gòu)

(1)卷積神經(jīng)網(wǎng)絡(luò)

首先，利用卷積層來對淺層提取的特征F0進(jìn)行進(jìn)一步的特征提取，得到卷積層提取的特征FC，如式(7)所示。

FC=fc(F0)

(7)

其中fc為FEB模塊中卷積層中特征提取的映射關(guān)系。

(2)輕量型Transformer網(wǎng)絡(luò)

輕量型Transformer是由2個(gè)STL(Swin Transformer Layer)模塊(如圖2(b)所示)串聯(lián)組成的，第一個(gè)STL模塊中MSA(Multi-headed Self-Attention)采用的是W-MSA(Window based Multi-headed Self-Attention)，W-MSA是窗口化的多頭自注意力機(jī)制，相較于傳統(tǒng)全局注意力機(jī)制減少了計(jì)算量。第二個(gè)STL模塊采用的則是SW-MSA(Shifted Window based Multi-headed Self-Attention)，由于W-MSA只能關(guān)注窗口本身的內(nèi)容，無法跨窗口連接，這就導(dǎo)致了窗口之間的特征信息無法傳遞，而SW-MSA可以通過引入移位窗口的方式，在保持窗口化的情況下兼顧了全局特征，并提高了計(jì)算效率，使得模型達(dá)到輕量化的效果。使用多層感知機(jī)(Multi-Layer Perception，MLP)，其中包括了兩個(gè)全連接層和GELU激活函數(shù)，來進(jìn)行進(jìn)一步的特征變化。在MSA和MLP之前都添加了LN(LayerNorm)層，并且這兩個(gè)模塊都引入了殘差連接。最后得到輕量化Transformer層提取的特征FT，如公式(8)所示。

FT=HSwin(F0)

(8)

其中HSwin為FEB模塊中的輕量型Transformer層。

(3)特征融合模塊

特征融合模塊含有多尺度卷積塊和特征重建部分，多尺度卷積塊利用不同尺寸的卷積核來進(jìn)行特征的多尺度提取，獲得不同感受野的特征。特征融合模塊采用殘差連接的方式，將多個(gè)維度的特征相加，從而使網(wǎng)絡(luò)更容易學(xué)習(xí)到低頻和高頻細(xì)節(jié)之間的映射關(guān)系，有助于加速訓(xùn)練過程。

在這個(gè)模塊中，將卷積層提取的特征FC和輕量型Transformer層提取的特征FT來進(jìn)行特征融合，得到FEB模塊的輸出特征FCf，如式(9)所示。

FCf=HCf(FC+FT)

(9)

其中HCf為特征融合模塊的函數(shù)。

1.2 動態(tài)調(diào)節(jié)模塊(DA)

圖3所展示的是動態(tài)調(diào)整模塊，主要運(yùn)用到的是全局平均池化(Global Average Pooling，GAP)操作。用于將卷積神經(jīng)網(wǎng)絡(luò)的特征圖轉(zhuǎn)換為固定長度的向量表示。GAP通過將每個(gè)通道的特征圖轉(zhuǎn)換為一個(gè)標(biāo)量值。具體來說，假設(shè)輸入圖像經(jīng)過卷積層后得到的特征圖為F，其尺寸為H×W×C，其中H和W是特征圖的高度和寬度，C是通道數(shù)，GAP會對每個(gè)通道的特征圖計(jì)算平均值，得到一個(gè)C維的向量表示，記為向量A，即每個(gè)通道的平均值。總而言之，通過GAP操作，特征圖的空間信息被壓縮成一個(gè)固定長度的向量。對向量A進(jìn)行一系列的線性變換和非線性激活操作，得到一個(gè)長度為C的通道權(quán)重向量S。將特征圖F與通道權(quán)重向量S進(jìn)行逐通道乘法操作，得到調(diào)整后的特征圖F′。最后再將調(diào)整后的特征圖F′輸入到后續(xù)的網(wǎng)絡(luò)層中，繼續(xù)進(jìn)行超分辨率重建的處理。

圖3 動態(tài)調(diào)整模塊作用流程圖

通過動態(tài)調(diào)整模塊，網(wǎng)絡(luò)可以根據(jù)輸入圖像的內(nèi)容和結(jié)構(gòu)，自適應(yīng)地增強(qiáng)或者減弱不同通道的影響，以提取到輸入圖像中最為主要的特征，同時(shí)減少網(wǎng)絡(luò)對無關(guān)或者冗余信息的依賴，從而改善超分辨率重建的質(zhì)量和細(xì)節(jié)保留的能力。

1.3 損失函數(shù)

網(wǎng)絡(luò)訓(xùn)練采用的損失函數(shù)是平均絕對誤差( Mean Absolute Error，MAE)，也稱為L1損失函數(shù)，如式(10)所示。

(10)

2 實(shí)驗(yàn)結(jié)果及分析

2.1 數(shù)據(jù)集和評價(jià)標(biāo)準(zhǔn)

在訓(xùn)練階段，采用DIV2K數(shù)據(jù)集來訓(xùn)練模型。該數(shù)據(jù)集包含有1 000張高清圖像(2K分辨率)，其中包括800張圖像作為訓(xùn)練數(shù)據(jù)，100張圖像作為驗(yàn)證數(shù)據(jù)，100張圖像作為測試數(shù)據(jù)。在評估階段，采用的基準(zhǔn)數(shù)據(jù)集是：Set5、Set14、BSD100、Urban100和Manga109。

在評價(jià)標(biāo)準(zhǔn)方面，采用峰值信噪比(PSNR)和結(jié)構(gòu)相似度(SSIM)來作為衡量模型效果的評價(jià)指標(biāo)。PSNR可以評價(jià)兩幅圖像之間的相似程度。PSNR值越高，說明重建出來的圖像中的失真或者誤差越小。較高的PSNR值表示重建圖像質(zhì)量較好，而數(shù)值較低意味著重建圖像存在著更明顯的偽影。SSIM與專注于對比像素差異的PSNR不同，SSIM主要比較圖像的亮度、對比度和結(jié)構(gòu)。通過計(jì)算三項(xiàng)的平均值來衡量相似性：亮度相似性、對比度相似性和結(jié)構(gòu)相似性。所以SSIM更加側(cè)重于圖像的結(jié)構(gòu)信息和感知質(zhì)量。

2.2 訓(xùn)練細(xì)節(jié)

網(wǎng)絡(luò)訓(xùn)練所用平臺為Ubuntu20.04，所有的實(shí)驗(yàn)均在單張 NVIDIA GeForce GTX 3060Ti顯卡上完成訓(xùn)練。在訓(xùn)練之前，對HR圖像進(jìn)行不同縮放因子的雙三次下采樣，生成對應(yīng)的LR圖像。本文利用L1損失和Adam算法對模型進(jìn)行優(yōu)化。初始學(xué)習(xí)率定義為1×10-4，算法中參數(shù)設(shè)置為β1=0.9、β2=0.999，訓(xùn)練周期定義為300K iterations，并且采用余弦退火衰減的學(xué)習(xí)方案來加速收斂，動量momentum為0.9。

2.3 模型比較

本文將STSR網(wǎng)絡(luò)與一些經(jīng)典的圖像超分網(wǎng)絡(luò)和近年來的輕量型圖像超分網(wǎng)絡(luò)進(jìn)行對比，其中包括FSRCNN[4]、VDSR[15]、MemNet[16]、EDSR[6]、CARN[8]、IMDN[17]、LatticeNet[18]、ESRT[19]。表1為在同一實(shí)驗(yàn)環(huán)境下所獲得測試結(jié)果。在每一行中，最好的結(jié)果用加粗的方式突出顯示。從測試結(jié)果可以看出在模型較小的情況下，量化指標(biāo)也取得了較好的結(jié)果，在模型性能和計(jì)算成本之間取得了較好的平衡。

表1 圖像超分辨率重建效果量化比較結(jié)果

2.4 消融實(shí)驗(yàn)

為了驗(yàn)證特征增強(qiáng)模塊FEB以及動態(tài)調(diào)整模塊DA的實(shí)際作用，本文將FEB模塊和DA模塊分別單獨(dú)引入到STSR模型中，通過實(shí)驗(yàn)來驗(yàn)證這兩個(gè)模塊在量化指標(biāo)PSNR和SSIM上的效果，可以驗(yàn)證出不同模塊對網(wǎng)絡(luò)性能提升的有效性。方法STSR-F是將STSR模型中的特征增強(qiáng)模塊FEB替換為單一的卷積神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)，由此來驗(yàn)證FEB模塊的作用。方法STSR-D是去掉STSR模型中的動態(tài)調(diào)整模塊DA，由此來驗(yàn)證DA模塊在網(wǎng)絡(luò)中的作用。

可以從表2中的量化指標(biāo)得出特征增強(qiáng)模塊FEB以及動態(tài)調(diào)整模塊DA對于本文網(wǎng)絡(luò)性能提升的有效性。

表2 消融實(shí)驗(yàn)

2.5 視覺效果評價(jià)

圖4中給出了STSR與其他圖像超分模型的視覺效果對比。本文設(shè)計(jì)模型重建出的高分辨率圖像包含了更多紋理細(xì)節(jié)，對于邊緣和線條等細(xì)節(jié)方面的重建，也展現(xiàn)出了更高質(zhì)量的重建效果。

圖4 STSR與其他模型的視覺對比

3 結(jié)論

針對基于混合構(gòu)架的圖像超分模型較高計(jì)算成本的問題，本文提出了一種基于卷積神經(jīng)網(wǎng)絡(luò)CNN和Transformer混合模型，結(jié)合了兩種構(gòu)架的優(yōu)勢，提高模型對局部細(xì)節(jié)特征和全局信息的建模能力，增強(qiáng)了上下文信息的利用效率，并且通過輕量化Transformer網(wǎng)絡(luò)，在達(dá)到較好重建效果的同時(shí)保持較低的計(jì)算成本。此外，本文還設(shè)計(jì)了DA模塊，根據(jù)輸入圖像來動態(tài)調(diào)整網(wǎng)絡(luò)的輸出，使重建的圖像更加貼近于真實(shí)圖像。量化指標(biāo)表明，本文方法在減少模型參數(shù)量的情況下，依然取得了較好的圖像重建效果。