基于自校準雙注意力的圖像超分重建網絡

周華平,郭潁杰,孫克雷

(安徽理工大學計算機科學與工程學院,安徽 淮南 232001)

給定一張低分辨率(Low Resolution,LR)圖片重建出對應的高分辨率(High Resolution,HR)圖片的過程,稱為單圖像超分辨率(Single Image Super Resolution,SISR)。重建后的圖像視覺質量較高,紋理細節豐富,這為克服醫學成像[1-2]、衛星成像[3]、安全與監控等各種計算機視覺應用中的分辨率限制提供了條件。得益于卷積神經網絡強大的特征表達能力,構建一個端到端的卷積神經網絡模型,學習LR圖片與HR圖片之間的非線性映射函數,成為目前主流的研究方法。

文獻[4]提出開創性工作SRCNN(super-resolution using deep convolutional networks),使用插值技術對LR圖像進行預處理,上采樣至目標尺寸作為模型的輸入,經過三層卷積提取圖像的特征。文獻[5]構建了一個極深的網絡結構VDSR(very deep convolutional networks),通過構建更深的網絡層獲得更大的感受野,并利用自適應梯度裁剪技術加快收斂過程。文獻[6]提出DRCN(deeply-recursive convolutional network),使用參數相同的遞歸塊來代替不同的卷積層,增加網絡深度同時降低參數量。文獻[7]提出深度持續記憶網絡(Deep Persistent Memory Network,Mem Net),利用記憶模塊多次遞歸,更好地利用多層特征信息。然而,上述模型通過對輸入特征進行放大處理,這會增加模型計算復雜度并產生可見的偽影。

文獻[8]提出的EDSR(enhanced deep residual networks)獲得了NTIRE2017 SR挑戰賽的冠軍,通過移除批歸一化層以及使用更深更寬的網絡結構來增加網絡表達能力。文獻[9]結合密集連接與殘差連接提出了RDN(residual dense network),充分利用層次結構特征。文獻[10]提出RCAN(residual channel attention networks),通過提高對通道間重要信息的關注度,使得網絡可以恢復更多的高頻信息,從而提升模型的重建效果。然而,上述算法為了獲得更好的重建性能,不斷加深網絡的層數,不可避免地導致參數量與計算量大幅度增加。此外,隨著網絡層數逐漸加深,高層次的特征更傾向于表達圖像中低頻的語義信息,導致對提高超分性能至關重要的紋理細節缺失。為了解決這類問題,提出了一種基于自校準雙注意力的圖像超分重建網絡(image super-resolution reconstruction network based on self-calibrated and dual-attention,SCDAN),主要貢獻如下:

(1)設計了一種自校準殘差塊(self-calibrated residual block,SCRB),使用校準分支和殘差塊分支分別提取特征信息。其中,校準分支使用像素注意力(pixel attention,PA)為特征圖中的所有像素點生成注意力系數,有目的地恢復圖像紋理信息。殘差塊分支利用跳躍連接將低頻信息與輸出特征相加,減緩梯度消失,穩定訓練。最后,對兩個分支的特征進行融合以提升模型的表達能力。

(2)在由多個SCRB構成的自校準殘差組(self-calibrated residual group,SCRG)的組內與組間進行分層特征融合,充分利用模型的淺層與深層信息,緩解由于網絡逐漸加深導致特征在傳播過程中信息缺失的問題。

(3)提出了一種雙注意力機制(dual-attention,DualA),結合通道注意力與空間注意力增強模型在通道和空間維度的建模能力,對融合后的信息進行過濾,保留圖像中的高頻細節。

1 自校準雙注意力圖像超分重建網絡

1.1 網絡整體結構概述

自校準雙注意力網絡(Self-calibrating Dual Attention Networks,SCDAN)主要通過級聯非線性映射模塊和圖像重建模塊實現高質量超分重建,其結構如圖1所示。定義網絡的輸入圖像為ILR,輸出圖像為ISR。

圖1 SCDAN整體網絡結構

網絡首先使用一個3×3大小的卷積層將輸入圖像ILR∈H×W×3(H、W分別表示圖像的高和寬)映射到高維空間,其操作可表示為

F0=f3×3(ILR)

(1)

式中,f3×3表示3×3卷積操作;F0表示淺層特征提取模塊的輸出特征。

接著將F0送入非線性映射模塊,進一步擬合輸入圖像ILR與輸出圖像ISR之間的非線性映射關系。

(2)

式中,fDualA表示雙注意力機制;f1×1表示1×1卷積操作;[·]表示特征通道級聯操作;Gi-1表示第i個SCRG輸入。則非線性映射模塊可表示為

FN=fDualA(f1×1([G1,G2,…,Gm]))

(3)

式中,FN表示非線性映射模塊的輸出特征。通過對SCRG的組內與組間引入分層特征融合結構,可以緩解隨著網絡加深,模型逐漸丟失部分紋理細節信息的問題,增強模型對特征的復用能力。融合后的特征可能包含有大量的冗余信息,通過設計的雙注意力機制對信息進行過濾,使得網絡更加關注特征通道與空間中的重要信息,可以恢復對圖像超分至關重要的高頻信息,獲得令人滿意的重建結果。

圖像重建模塊將淺層特征提取模塊的輸出F0與非線性映射模塊的輸出FN相加作為輸入特征,通過亞像素卷積層進行上采樣,再經過一個3×3卷積層調整通道維度得到輸出圖像ISR。其操作可表示為

ISR=f3×3(fPixelShuffle(F0+FN))

(4)

式中,fPixelShuffle表示亞像素卷積操作。

1.2 自校準殘差塊

LR圖像中含有大量易于學習的低頻信息,利用殘差學習可以將這些信息通過恒等映射直接傳輸到網絡末端,使網絡直接學習對圖像超分任務至關重要的高頻信息,緩解梯度消失問題。文獻[11]提出一種自校準卷積,將輸入特征進行通道拆分,以并行的雙分支結構對特征進行特征提取和重校準。結合殘差塊與自校準卷積的優點,提出了一種自校準殘差塊(SCRB),利用傳統的殘差塊作為其中一個分支,用于捕獲LR圖像的特征信息。另一分支使用像素注意力(PA)為特征圖所有像素點生成注意力系數,有目的地恢復圖像的紋理信息。

如圖2所示,設計的自校準殘差塊由并行的殘差塊分支和校準分支構成。輸入特征Xn-1在經過兩條支路前,會經過兩個1×1卷積層進行降維,將通道數減半以減少模型參數量,同時得到兩條支路的輸入H1和H2。其操作可表示為:

圖2 自校準殘差塊

(5)

(6)

H1,H2首先分別經過殘差塊分支和校準分支,然后對兩個分支的輸出特征進行特征融合并與輸入特征Xn-1相加,得到輸出特征Xn。其操作可表示為:

(7)

F2=f3×3(H2)⊙σ(f1×1(H2))

(8)

Xn=f1×1([F1+F2])+Xn-1

(9)

式中,Φ表示LReLU激活函數;σ表示Sigmoid激活函數;⊙表示逐元素相乘操作;F1,F2分別為殘差塊分支和校準分支的輸出。

1.3 雙注意力機制

通過分層特征融合操作,可以充分復用各層級的輸出特征,但是會帶來大量冗余信息。為了提高網絡提取圖像高頻細節的能力并減少冗余信息對后面網絡層的影響,SCDAN在組內與組間的分層特征融合操作后加入一個雙注意力機制(DualA),對融合后的特征進行信息過濾,使得網絡更加關注通道與空間的重要特征,忽略那些不重要的特征。

DualA結構如圖3(a)所示,主要采用通道注意力(channel attention,CA)和空間注意力(spatial attention,SA)并行的模式分別處理通道維度和空間維度上的重要信息。Transformer中廣泛使用的多頭自注意力機制可以建立像素間的長程依賴關系,但是其計算量隨輸入大小呈二次增長。受到文獻[12]啟發,通過在特征的通道維度而非空間維度上計算像素間的相關性,可以降低模型的計算負載,計算量可從O(W2H2)降至O(C2)。對于通道注意力分支(CA),其結構如圖3(b)所示。給定輸入特征Xn-1,CA分支首先使用3個1×1卷積生成查詢向量(Q)、鍵向量(K)和值向量(V)。其操作可描述為:

圖3 雙注意力機制組件

(10)

(11)

(12)

通過reshape操作調整向量的維度得到新的Q′∈C×HW、K′∈C×HW和V′∈C×HW。然后,將Q′與K′的轉置執行點乘操作并經過Softmax函數生成權值圖A∈C×C。最后,將權值圖與V′相乘,得到加權后的特征FCA,其操作可表示為

FCA=Softmax(Q′·(K′)T/α)·V′

(13)

式中,Softmax表示用于生成權值圖的Softmax函數;α是一個可學習的縮放參數,用于控制Q′和K′的點積的大小。通過對特征通道間的長程依賴關系進行建模,可以使網絡能夠捕獲遠距離相似紋理,提升模型的特征表達能力。

為了進一步恢復圖像精確的空間結構,空間注意力分支(SA)首先對輸入特征在通道維度上使用平均池化和最大池化生成兩個大小為1×H×W的特征圖,然后對兩個特征圖按通道進行拼接,并通過卷積核大小為7×7的卷積層提取特征圖的空間信息,最后經過Sigmoid函數生成權值圖并對輸入特征進行加權,其結構如圖3(c)所示。該過程可表示為

FSA=Sigmoid(f7×7([AvgPool(Xn-1);MaxPool(Xn-1)]))⊙Xn-1

(14)

式中,AvgPool(·)和MaxPool(·)分別表示平均池化和最大池化操作;f7×7表示7×7卷積;FSA表示SA分支的輸出。將兩個分支的輸出特征圖進行融合并通過1×1卷積進一步調整,得到輸出特征Xn,使得模型能夠更好地恢復圖像的高頻信息。其操作可表示為

Xn=f1×1(FCA(Xn-1)+FSA(Xn-1))

(15)

通過結合通道注意力分支與空間注意力分支,增強了模型在通道和空間維度的建模能力,提升了模型的判別能力,使得模型重建結果更加清晰,紋理細節更加豐富。

1.4 損失函數

(16)

式中,fSCDAN為本文提出的網絡模型;θ為網絡中的可訓練參數。

2 實驗

2.1 數據集與評價指標

本文實驗在公共數據集DIV2K[13]數據集和Flickr2K[14]上進行訓練,共包含3 450張對應的低分辨率圖片(LR)和高分辨率圖片(HR)。其中,在DIV2K數據集中選取前800張和Flickr2K組成實驗訓練集。實驗測試集包含:Set5[15]、Set14[16]、BSD100[17]、Urban100[18]。本文將峰值信噪比(PSNR)和結構相似度(SSIM)作為性能評價指標,并在YCbCr通道的Y通道(亮度)進行計算。

2.2 參數設置

在模型訓練時采用Adam算法對模型進行優化。批大小設置為16。數據增強使用隨機旋轉90°、180°、270°和水平翻轉。本文對2×,3×和4×三種縮放因子進行實驗,將大小為48×48的LR子圖作為模型的輸入,對應的真實標簽HR在3種縮放因子下的大小為96×96、144×144和192×192。實驗初始時,學習率為5×10-4,每2×105個迭代學習率減半,總迭代次數為8×105。模型使用的深度學習框架為Pytorch,編程語言為Python3.8,在NVIDIA RTX 3090進行訓練和測試。

2.3 實驗結果與分析

為了驗證本文提出結構的有效性,進行了一系實驗。為了公平比較,所有實驗在同一環境下,迭代次數為5×105。

1)自校準殘差塊有效性分析 自校準殘差塊(SCRB)是非線性映射模塊的核心結構,為了驗證SCRB的有效性,本文與傳統的殘差塊(ResBlock)結構進行了對比實驗,結果如表1所示。其中,“ResBlock”表示用傳統的殘差塊替代SCRB。相比于傳統的殘差塊,使用SCRB將PSNR平均提高了0.03dB,而參數量與計算量大大降低。這得益于校準分支使用像素注意力(PA)為特征圖所有像素點生成了權重信息,可以有目的地恢復圖像紋理信息。

表1 自校準殘差塊(SCRB)性能分析

2)特征融合與雙注意力機制有效性分析 為了證明設計的雙注意力機制與特征融合結構的有效性,進行了5組消融實驗,結果如表2所示。其中,“FF”表示特征融合,“CA”為通道注意力分支,“SA”為空間注意力分支。由表2中前2組消融實驗對比分析可得,在自校準殘差組的組內與組間進行分層特征融合在4個測試集上的PSNR平均提升了0.047 5dB。這表明通過聚合淺層與深層特征,可以緩解因網絡不斷加深導致的信息丟失,充分利用各層特征,提升模型的表達能力。

表2 特征融合與雙注意力機制性能分析

表2中后3組實驗主要對雙注意力機制的有效性進行評估。通過表2結果分析,在使用分層特征融合的基礎上,單獨使用通道注意力分支和空間注意力分支分別將4個測試集上的PSNR平均提升了0.075dB和0.06dB。這表明使用注意力機制可以對聚合后的信息進行過濾,更加關注那些對于圖像恢復至關重要的高頻信息。雙注意力機制同時關注特征的通道和空間維度的重要特征,在4個測試集上的PSNR性能指標平均提高了0.15dB,超越了單獨使用通道注意力或空間注意力,充分體現了本文設計的雙注意力機制在處理圖像高頻信息的有效性。

3)復雜度分析 本文網絡在Set5數據集4×放大倍數下的PSNR與不同網絡的參數量(Params)與浮點量(Flops)對比如表3所示,比較網絡包括VDSR、DRCN、CARN[19]、LAPAR-A[20]、IMDN、LBNet[22]。對于輕量化網絡來說,可訓練參數的數量是需要考慮的基本元素。VDSR雖然參數量要小于本文算法,但是該算法的輸入與輸出大小相等,導致模型計算量急劇增加。相較于DRCN、CARN,本文模型的參數量僅為這兩種模型的49%和55%,但PSNR卻分別提高了0.9dB和0.3dB。IMDN與LBNet雖然參數量與計算量略小于本文算法,但是本文算法的重建性能更有優勢。圖4對比了本文網絡與其他輕量級網絡在Set5數據集2倍放大上的參數量與PSNR值的對應關系,本文所提網絡保持較低參數量的同時,取得了最好的性能。這表明本文所提網絡可以在性能和參數量上取得很好的權衡。

表3 4×放大倍數的復雜度分析

圖4 Set5數據集2×放大倍數模型參數量和PSNR值

2.3 與其他輕量級網絡的比較

1)定量客觀比較 將本文所提網絡與其他輕量級網絡在2×、3×、4×放大倍數上進行了定量客觀比較,包括:SRCNN、VDSR、DRRN[23]、DRCN、MemNet、IDN[24]、CARN、IMDN、LAPAR-A、RFDN[25]、LBNet、SCET[26]。實驗對比結果如表4所示,所提網絡在4個測試集上PSNR和SSIM性能指標均達到最優,尤其在Urban100數據集上性能提升較為明顯,該數據集中主要含有大量的城市圖片,包含大量的邊緣信息且細節分布在不同頻段,對于圖像超分任務來說十分具有挑戰性。本文網絡在Urban100數據集2×、3×、4×放大倍數對比參數量最多的DRCN,PSNR分別提高1.64dB、1.31dB、1.21dB。這主要歸因于提出的雙注意力機制,其中,CA分支利用自注意力機制對全局相關性進行建模,建立圖片中像素間的遠距離相關性,而城市建筑多為重復性結構,所以可以更好地恢復遠距離相似紋理特征。同時,SA可以捕獲圖像中的空間位置信息,可以更好地恢復復雜的結構特征。

表4 在四種數據集上,不同網絡在2×、3×和4×上的平均PSNR和SSIM定量比較

2)主觀效果展示與分析 本節詳細分析了所提網絡與其他輕量級網絡的可視化結果,包括:Bicubic、SRCNN、IDN、LAPAR-A、CARN和IMDN。圖5展示了在放大倍數為4時,不同網絡對Set14數據集中的圖片“Barbara”的重建效果?？梢钥闯?其他網絡恢復的書籍邊緣均出現了扭曲、重疊,本文網絡重建的書籍邊緣更加銳利。圖6展示了在放大倍數為4時,各種網絡對BSD100數據集中“253027”的重建效果。其他網絡重建的斑馬條紋較為模糊,而本文所提網絡重建效果較為清晰且恢復了更多的高頻細節。

圖5 放大倍數為4時,不同網絡對Set14數據集中的“barbara”圖像局部放大效果可視化對比

圖7、圖8展示了在放大倍數為4時,本文網絡與其他算法對Urban100數據集中的兩張圖片“img_005”和“img_096”的重建效果。雙三次插值、SRCNN生成結果比較模糊,失去了大部分結構信息。IDN、LAPAR-A、和IMDN對于墻體的邊緣恢復較差,出現不同程度的模糊、扭曲,無法恢復銳利的紋理細節。相比之下,本文網絡可以恢復更多的細節和更清晰的邊緣,并獲得更高的視覺質量,這歸功于更有效的特征提取和特征復用。

圖7 放大倍數為4時,不同網絡對Urban100數據集中的“img_005”圖像局部放大效果可視化對比

圖8 放大倍數為4時,不同網絡對Urban100數據集中的“img_096”圖像局部放大效果可視化對比

4 結論

在保持模型性能的基礎上減少模型的參數量與計算量,是SR領域熱門研究方向。本文提出了一種基于自校準雙注意力的圖像超分重建網絡,設計了一種自校準殘差塊,利用自校準的思想有目的地恢復圖像紋理信息。將特征融合結構與設計的雙注意力機制相結合,充分復用各層級的輸出特征,同時過濾冗余特征,更加關注融合后的特征在通道和空間維度的重要信息,以恢復圖像缺失的高頻信息。在4個測試集上的定量與定性分析表明,與其他輕量級方法相比,模型在視覺質量和客觀指標方面均取得不錯地結果。同時,所提算法在模型性能與復雜度之間取得了很好的平衡。