基于雙注意力機制的輕量級圖像超分辨率重建

馮興杰 王 榮

1(中國民航大學信息網絡中心 天津 300300) 2(中國民航大學計算機科學與技術學院 天津 300300)

0 引 言

在實際應用中,受到圖像采集設備成本、傳輸帶寬、存儲年限久遠等問題的限制,很難保證得到的圖像均為精致的高分辨率圖像。因此超分辨率重建(Super-Resolution,SR)技術的產生顯得尤為重要。根據輸入數據類型的不同,可以將SR分為單圖像SR和多圖像SR,由于多圖像SR直接基于單圖像SR,因此本文主要討論單圖像超分辨率重建(Single Image Super-Resolution,SISR)。

圖像超分辨率重建的目的是將給定的具有粗糙細節的低分辨率圖像(LR)轉換為相應的具有更好的視覺效果以及精致細節的高分辨率圖像(HR)。SR的應用領域很廣泛,如目標檢測(尤其是針對小對象)、監控視頻、衛星遙感等。實際上,無限多的高分辨率圖像可以通過下采樣得到相同的低分辨率圖像。因此,SR問題本質上是不適定的,不存在唯一解。為了緩解這一問題,文獻中提出了許多SISR方法,目前的SISR方法一般可分為兩類:基于重建的方法(Reconstruction-Based)[1-2]和基于樣本學習的方法(Example Learning-Based)。基于重建的方法僅用了一些先驗知識來正則化重建過程,計算量低且簡單,不過無法處理復雜圖像結構。基于學習的方法利用大量通過學習得到的先驗數據來學習低分辨率圖像到高分辨圖像的某種映射關系,用學習到的映射關系來預測高分辨率圖像,常見的有基于稀疏編碼和基于深度學習的方法。

近年來深度學習發展迅猛。Dong等[3]首先提出了一個采用三層卷積神經網絡的SRCNN模型對特征提取、非線性映射和圖像重建階段進行端到端的聯合優化。Shi等[4]提出了一種有效的亞像素卷積神經網絡ESPCN,它直接在LR空間中提取特征映射,有效地降低了計算復雜度。受到功能強大的深度卷積神經網絡(Convolutional Neural Networks,CNN)的影響,許多基于CNN的SR方法嘗試訓練一個更深的網絡以獲得更好的重建性能。Kim等[5]提出了一種包含20個卷積層的VDSR模型,該模型采用殘差學習和自適應梯度裁剪來減輕訓練難度,證明了更深的網絡可以提供更好的特征重建,使網絡捕捉到更加全局的特征,進而有益于超分辨率重建工作中圖像細節的恢復。

因此在之后提出的一些模型[6-9]為了獲得更好的性能,不斷深化或拓寬網絡已經成為了一種趨勢。然而,不斷對網絡進行深化和拓寬,隨之而來的則是實際應用中算法復雜度和內存消耗的提升。在實際應用中,如移動和嵌入式視覺應用中,這些方法的適用性較差。此外,傳統的卷積層平等地對待圖像的高頻特征與低頻特征,不加區分地將特征傳向后續層。但高頻特征往往包含更多的細節信息,而SR任務的重點是恢復圖片中的紋理細節信息,因此應該將重點放在擁有高頻細節的特征上。隨著研究的深入,注意力的重要性得到了廣泛的關注,注意力不僅闡明了關注的焦點在哪里,還提高了焦點的表達。與被廣泛關注的Hu等[10]提出的只關注通道注意力的方法不同,Woo等[11]對基于通道和空間兩個維度上的特性響應進行了研究,提出了通道和空間注意力機制。它可以自適應地重新校準通道和空間兩個維度上的特性響應以提高網絡的表達能力,很好地解決了深層網絡中特征不加區分的向后傳遞的問題,同時該注意力機制引入的參數量很少。這為SR獲取全局特征來更好地恢復高頻紋理細節以及網絡的輕量化提供了新思路。

基于上述討論,本文提出一種基于雙注意力機制的輕量級圖像超分辨率重建(SR-LAM)算法。所提出的SR-LAM網絡為了能夠充分地捕獲上下文信息,將網絡深度保持在20層。網絡結構分為四部分。第一部分由一個增強模塊對LR圖像進行淺層特征的提取及校準。第二部分由包含組卷積的卷積核大小為3×3的特征提取塊堆疊,逐步細化殘差信息。第三部分與第一部分相同,用一個增強模塊進行深層特征提取和重新校準。最后由重建模塊將殘差信息上采樣到目標尺寸后與上采樣到目標尺寸的輸入圖像融合輸出結果。最后通過在縱向和橫向上進行對比實驗,表明了盡管使用較少的卷積層,本網絡仍然能夠獲得具有競爭性的結果。

1 方法與理論

自Kim等[5]的方法被提出以來,許多超分辨率重建算法都是基于它進一步優化網絡結構以提高網絡性能。本文也將在VDSR模型的基礎上提出改進。

該模型以插值后的目標尺寸低分辨率圖像作為網絡的輸入,這無疑加大了網絡的計算量。為了擁有更大的感受野讓圖像特征提取到更多的細節,VDSR采用了更深(20層)的網絡結構,為了加速收斂,采用了非常高的學習率,使用殘差學習來緩解高學習率導致的梯度問題。Kim等通過實驗證明了不斷加深網絡會帶來感受野的增加,使網絡能夠捕捉到更加全局的特征,進而有益于SR工作中紋理細節的恢復。

雖然Kim等對超分辨重建算法做出了很好的改進,但其中仍存在著一些影響超分辨率重建質量和效率的因素:

1) 對于目標尺度系數n,與插值后的目標尺寸LR圖像進行卷積的計算量將是與原始LR圖像進行卷積計算的n2倍,另外,雙三次插值方法不會帶來額外的有利于解決不適定問題的信息。

2) VDSR及之后提出的一些模型為了獲得更好的性能,不斷深化或拓寬網絡已經成為了一種趨勢,但隨之而來的是龐大的計算量及網絡訓練困難。

3) LR圖像與HR圖像之間存在許多相似的部分,SR任務的重點則是恢復圖片中的紋理細節信息,因此應該將重點放在擁有高頻細節的特征上。VDSR及一些傳統的卷積網絡認為所有空間位置和通道對超分辨率具有統一的重要性,通常采用級聯網絡拓撲,每一層的特征映射被無區別地發送到后續層,這不利于把注意力集中在更有利于恢復高頻細節的特征上。

針對上述現象,隨著研究的深入,本文對其做出了改進:

1) 本文方法采用后上采樣代替提前上采樣,以此來降低網絡運算量,這有利于網絡輕量化的實現。

2) 在某些情況下,考慮到并非所有特征都是SR所必需的,其重要性各不相同,有選擇地側重于某些特定特征對更好地恢復高頻細節是有幫助的。由于卷積運算是通過混合跨通道和空間信息來提取特征的,所以強調沿著通道和空間軸這兩個主要維度的重要特征是有必要的。對空間和通道兩個維度上的特征響應研究[11]允許這種靈活性的產生。提出雙注意力機制可以自適應地重新校準通道和空間兩個維度上的特性響應,提高網絡的表達能力,解決深層網絡中特征不加區分的向后傳遞的問題。另外該注意力機制引入的參數量非常少。融合高效且輕量的雙注意力機制代替普通卷積塊在保證了網絡質量的情況下進一步實現了輕量化,很好地解決了網絡加深帶來的訓練困難。

3) 用組卷積代替普通卷積已經被證明[12]了能夠在性能損失很小的前提下減少大量的網絡參數,本文方法為了進一步實現網絡輕量級在特征提取模塊融入了組卷積。

1.1 模型結構

圖1展示了本模型的總體網絡架構。模型由四部分組成,包括三種模塊:增強模塊、特征提取模塊、重建模塊。

圖1 SR-LAM網絡結構

SR的具體任務是將原始高分辨率圖像(IHR)進行下采樣得到輸入網絡的低分辨率圖像(ILR),經網絡得到最終的重建后高分辨率圖像(ISR)。其中的下采樣操作其實是在模擬原始圖像在其獲取、存儲和傳輸過程中存在不可避免的各種圖像退化過程,退化函數通常是未知的,在目前的SR方法中大多采用雙三次插值對IHR進行下采樣預處理來模擬這個過程,為了與其他網絡保持一致,本文也采用此插值方法來處理原始高分辨率圖像。

對于網絡的輸入圖像,第一部分由一個增強模塊對LR圖像進行淺層特征的提取及校準。其中增強模塊包含一個卷積核大小為3×3的卷積層和基于通道和空間的雙注意力模塊。第二部分由5個包含組卷積的卷積核大小為3×3的特征提取塊堆疊,逐步提取殘差信息。第三部分與第一部分相同,用一個增強模塊進行深層特征提取和重新校準。最后由重建模塊將信息采用[4]提出的亞像素卷積的方式上采樣到目標尺寸并和直接雙三次插值上采樣后的LR圖像融合相加。

1.2 增強模塊

增強模塊由一個3×3的卷積和一個基于通道和空間兩個維度的雙注意力模塊組成。卷積層后面跟著一個LReLU激活函數,如圖2所示。首先進行淺層特征的提取,緊接著對特征進行一次校準,讓隨后的特征提取步驟把注意力集中在更有利于恢復高頻細節的特征上。輸出維度為64的特征圖。用公式表示如下:

圖2 增強模塊

Fout=FS(FC(H(Fin)))

(1)

式中:Fin表示輸入;H表示特征提取函數;FC表示通道注意力操作;FS表示空間注意力操作;Fout為增強模塊的輸出。

此注意力模塊是一個輕量級模塊,在僅引入了1 194個參數的情況下提高了網絡的表達能力,滿足了對網絡輕量級的要求。該模塊將注意力的特征細化應用于通道和空間兩個不同維度,加強了對高頻紋理細節的校準能力,引入該模型使得本網絡可以在保持較小開銷的同時獲得有競爭力的重建效果。

在通道注意力子單元中,如圖2所示。對于輸入的特征圖F∈RH×W×C首先使用平均池化和最大池化來聚合特征圖的空間信息,生成兩個不同的空間上下文信息來分別表示平均池化特征和最大池化特征。接著將兩個上下文特征分別發送到共享網絡,共享網絡由含有一個隱藏層的MLP組成。之后對MLP輸出的兩個特征進行逐元素求和以及Sigmoid激活操作,生成通道注意力特征圖AC(F),該特征圖和輸入特征圖逐元素相乘得到空間注意力子模塊的輸入特征圖FC,這個過程可以表示為:

FC=sigmoid(MLP(Avgpool(F))+MLP(Maxpool(F)))

(2)

在空間注意力子單元中,如圖2所示。通道注意力子模塊的輸出Fc即本子模塊的輸入。首先對其沿著通道軸進行平均池化和最大池化得到兩個上下文信息,將得到的兩個信息基于通道方向做拼接,經一個卷積降維到單通道,然后經Sigmoid激活函數生成最終的空間注意力特征圖。最后將該特征圖與輸入特征圖做乘法得到FS,這個過程可以表示為:

FS=sigmoid(f7×7(Avgpool(FC);Macpool(FC))

(3)

式中:f7×7為一個7×7×1的卷積層。

上述過程可以簡潔地概括為:

FC=AC(F)?F

(4)

Fout=FS=AS(FC)?FC

(5)

式中:AC為通道注意力操作;AS為空間注意力操作。

1.3 特征提取

用組卷積代替普通卷積已經被證明[12]可以減少大量的參數和操作,而性能損失很小。為了進一步降低網絡參數,減少網絡計算成本。本文在特征提取模塊的第二層引入了組卷積。組卷積可看成是一種稀疏卷積連接的形式。將輸入特征在通道方向上分成若干組,分別對每一組進行卷積后在通道方向進行拼接。如圖3所示。特征提取模塊由兩個3×3的卷積層和一個組卷積塊組成。使用組卷積的優勢在于它使模型的效率可調,用戶可適當地選擇組大小,因為組大小和性能處于權衡關系中。本文延續Ahn等[12]所提出結論,設定組卷積的組大小為4。

圖3 特征提取模塊

1.4 重建模塊

早期的上采樣方法大多是采用雙三次插值進行提前上采樣,即在網絡的第一層之前或者第一層對輸入的LR上采樣到與輸出的HR相匹配的尺寸。這種方法顯然會增加計算復雜度,尤其是對于卷積操作,因為卷積網絡的處理速度直接取決于輸入圖像的分辨率,這不符合本文輕量級網絡的目標。其次,雙三次插值方法不會帶來額外的有利于解決不適定問題的信息。因此本文沒有延續VDSR中的提前上采樣操作,而是采用了文獻[4]中提出的亞像素卷積進行后上采樣。亞像素卷積通過通道擴增和像素點重排來實現圖像放大。由于輸入的圖片分辨率降低,可以有效地使用較小的卷積核來進行特征提取,同時維持給定的上下文區域。分辨率和卷積核尺寸的減小也大大降低了計算量和內存的復雜度。

最后將重建模塊的輸出與經上采樣到目標尺寸的輸入圖像融合得到最終結果。

1.5 損失函數

本文采用峰值信噪比(PSNR)和結構化相似度(SSIM)作為重建的評價指標。PSNR與還包含L1和L2的像素損失高度相關,像素損失最小直接最大化PSNR。本文采用L1作為模型的損失函數,L1公式如下:

(6)

2 實驗訓練及結果

2.1 數據集

該模型使用DIV2K數據集[13]進行訓練,它包括1 000幅圖像,其中:800幅訓練圖像;100幅驗證圖像;100幅測試圖像。它是一種新提出的高質量圖像數據集。由于該數據集的豐富性,最近的SR模型[14-16]也使用了DIV2K數據集。本文方法在四個被廣泛使用的標準數據集上進行評估:Set5、Set14、BSD100、Urban100。

2.2 實驗細節

本文使用文獻[17]中提出的Xavier初始化方法進行權重初始化,偏置項初始化為零。使用Adam[18]對網絡進行優化。初始學習率為:1E-4,最大迭代次數為10 000,每2 000次迭代學習率減半。

本文的實驗環境如表1所示。

表1 實驗環境

2.3 實驗結果及分析

為了探索通道和空間的注意力機制對SR任務的作用,首先將本文模型與剔除掉注意力模塊的NA-SR-LAM進行了對比。訓練過程如圖4所示,橫軸為訓練輪數,縱軸為PSNR值。可以看到包含注意力模塊的模型訓練過程更穩定,隨著訓練次數的增加,波動逐漸變小。實驗結果如表2所示。實驗結果顯示迭代次數為10 000時,添加了通道和空間注意力的模型較沒有添加的模型在四個測試數據集上的峰值信噪比值分別提高了:0.29 dB,0.27 dB,0.24 dB,0.53 dB。此外,一個普通的3×3卷積核將引入36 928個參數,一個注意力模塊僅引入了1 194個參數。綜上所述,通道和空間注意力模塊不僅有利于增強特征表達能力而且有利于模型的輕量化。

表2 SR-LAM與無注意力的SR-LAM對比

圖4 SR-LAM與無注意力的SR-LAM對比(S=2)

表3展示了迭代次數為3 000時,Set14數據集上有無組卷積情況下模型參數量以及PSNR的對比。其中N-Group表示用普通卷積來代替特征提取模塊中組卷積。

表3 SR-LAM與無組卷積的SR-LAM對比

由表3可知,融合組卷積的模型與未應用組卷積的模型在PSNR上僅相差0.01 dB,但參數量卻比未應用組卷積的模型降低了將近23%,因此引入組卷積是實現該模型輕量化的有效措施。

本文提出的改進策略在不同程度上都對網絡的輕量化做出了一定貢獻,為了進一步直觀展示本文算法是否更好地平衡了模型復雜度與重建效果,本文對近年來基于深度學習的經典超分辨率算法進行了對比實驗。實驗條件:目標尺度系數(S)為4,數據集為Set14。

選取的對比算法有:SRCNN[3]、LapSRN[19]、FSRCNN[20]、VDSR[5]、DRRN[6]、IDN[16]、本文提出的SR-LAM算法:SRCNN模型采用三層卷積神經網絡對特征提取、非線性映射和圖像重建階段進行端到端的聯合優化;FSRCNN方法將SRCNN的提前上采樣改為了后上采樣,并且用8個小尺寸的卷積核來代替之前的大卷積核,降低了計算復雜度,重建效果略有提升;VDSR將網絡加深到20層以提高圖像重建效果;DRRN延續了VDSR更深的網絡層次重建效果更好的思想,將網絡加深到52層,重建效果略有提升,模型參數量也有所下降;LapSRN方法對原始LR圖像先生成低倍放大圖像,再逐步細化生成高倍放大的圖像,很好地改善了放大倍數高的復雜度問題;IDN直接從LR圖像提取特征,減少了模型計算量,受到注意力思想的影響,作者提出了可提高網絡表達能力、壓縮特征冗余信息的蒸餾模塊,通過對此模塊的疊加在一定程度上提高了網絡重建效果。

實驗結果如圖5所示,其中橫軸表示算法所需參數量,縱軸為重建效果評價指標PSNR,不同大小的圓點代表不同算法所需參數量。結果表明:本文提出的方法更好地兼顧了重建效果與模型復雜性。與重建效果相當的IDN、LapSRN相比,本文所提出的模型更加簡潔,參數量更少。在模型參數量方面,SR-LAM的重建效果遠好于參數量相當的DRRN。對比基礎模型VDSR,其參數量是本文所提方法的1.5倍。

圖5 不同算法在Set14數據集上的PSNR及參數個數對比(S=4)

圖6直觀地對比了目標重建系數為2的comic圖像的重建效果。對比Bicubic、SRCNN、VDSR、IDN方法,本文方法視覺上手指邊緣以及流水的線條都更加清晰,這是因為本文將更多的注意力放在了高頻信息,因此本方法對線條以及高頻信息的重建效果更好、紋理更加清晰。

圖6 不同方法對comic重建的視覺效果對比(S=2)

表4顯示了最近基于CNN的SR方法的性能和模型大小。選取的算法包括:LapSRN[19]、IDN[16]、RCAN[21]、SAN[22]。本文提出的SR-LAM較LapSRN和IDN擁有更少的參數和更好的性能。雖然RCAN和SAN性能略勝,但其付出的代價特別大,二者平均參數量是本文所提方法的32倍。

表4 參數及性能比較(Set14)

表5客觀地對比了本文方法與其他五種經典方法(雙三次插值(Bicubic)、SRCNN、FSRCNN、VDSR、IDN)在不同上采樣倍數下的PSNR以及SSIM。

表5 本文方法(SR-LAM)與其他方法結果比較

結果顯示本文模型取得了具有競爭力的結果:與基礎模型VDSR相比,在S分別等于2、3、4的情況下,四個測試數據集上的PSNR值平均提高了0.273 dB、0.253 dB、0.295 dB,SSIM值平均提高了0.003、0.005、0.008;與先進的IDN模型相比,在Set14和Urban100數據集上本文方法取得了更好的結果,而且通過對比圖6(e)和圖6(f)可以發現,本文方法對手指及指甲處的邊緣恢復得更為清晰,視覺效果更好,與原圖像也更加相似。

3 結 語

本文提出一個簡潔、輕量、便于在移動端使用的超分辨率重建模型:SR-LAM。方法上主要通過:融合高效且輕量的雙注意力機制代替普通卷積塊、采用后上采樣方法來代替提前上采樣、引入部分組卷積的方式,使模型能夠在擁有少量網絡參數的情況下達到一個有競爭力的重建效果。實驗結果證明,對比重建效果相當的先進模型,本文方法僅使用了其1/2,甚至更少的參數量。因此可以說本文模型更好地平衡了模型復雜度與重建效果。后續將進一步研究如何將這種高效、簡潔的網絡用于視頻超分辨率領域。