CVPR 2022 | Restormer: 刷新多個(gè)low-level任務(wù)指標(biāo)

arXiv：https://arXiv.org/abs/2111.09881

code：https://github.com/swz30/Restormer

本文是MPRNet與MIRNet的作者在圖像復(fù)原領(lǐng)域的又一力作，也是Transformer技術(shù)在low-level領(lǐng)域的又一個(gè)SOTA。針對Transformer在高分辨率圖像復(fù)原中存在的難點(diǎn)，提出了兩種MDTA與GDFN兩種改進(jìn)，極大程度上緩解了計(jì)算量與GPU緩存占用問題。所提方案刷新了多個(gè)圖像復(fù)原任務(wù)的SOTA性能。

Abstract

通過MHSA與FFN進(jìn)行改進(jìn)，本文提出一種高效Transformer，它可以捕獲長距離像素相關(guān)性，同時(shí)可適用于大尺寸圖像。所提方案Restormer(Restoration Transformer)在多個(gè)圖像復(fù)原任務(wù)上取得了SOTA性能，包含圖像去雨、圖像去運(yùn)動(dòng)模糊、圖像去散焦模糊以及圖像降噪(包含合成與真實(shí)噪聲)，可參見下圖。

本文主要貢獻(xiàn)包含以下幾點(diǎn)：

• 提出了一種編解碼Transformer用于高分辨率圖像上多尺度local-global表達(dá)學(xué)習(xí)，且無需進(jìn)行局部窗口拆分；
• 提出一種MDTA(Multi-Dconv head Transposed Attention)模塊，它有助于進(jìn)行局部與非局部相關(guān)像素聚合，可以高效的進(jìn)行高分辨率圖像處理；
• 提出一種GDFN(Gated-Dconv Feed-forward Network)模塊，它可以執(zhí)行可控特征變換，即抑制低信息特征，僅保留有用信息。

Method

本文旨在設(shè)計(jì)一種高效Transformer模型，它可以處理復(fù)原任務(wù)中的高分辨率圖像。為緩解計(jì)算瓶頸，我們對MHSA進(jìn)行了關(guān)鍵性改進(jìn)并引入多尺度分層模塊，見下圖。

Overall Pipeline ? 給定退化圖像 , Restormer的處理流程如下:

• 首先, 采用卷積得到底層特征 ;
• 然后, 這些底層特征將經(jīng)由 4 級對稱編解碼器處理變換為深層特征 (每級編解碼器包含 多個(gè)Transformer模塊, 編碼器分層降低空間分辨率, 同時(shí)擴(kuò)展通道, 解碼器則降低通道數(shù), 提升空間分 辨率), 注: 為輔助重建過程, 編解特征將與解碼特征通過跳過連接進(jìn)行concat并后接 卷積降維;
• 其次, 深層特征將經(jīng)由提煉模塊在進(jìn)一步豐富特征多樣性;
• 最后, 通過卷積層處理生成殘差圖像 , 與輸入圖像相加即可得到復(fù)原圖像 。

Multi-Dconv Head Transposed Attention

Transformer的主要計(jì)算負(fù)擔(dān)來自自注意力層。在常規(guī)自注意力中，key-query的點(diǎn)乘操作會(huì)隨分辨率以增長。因此，基本不可能將自注意力層用于圖像復(fù)原任務(wù)(圖像復(fù)原的輸入分辨率一般都比較大)。

為解決該問題，我們提出了MDTA(見上圖)，它具有線性復(fù)雜度，其關(guān)鍵成分在于：在通道維度(而非空間維度)執(zhí)行自注意力計(jì)算跨通道的交叉協(xié)方差以生成關(guān)于全局上下文的隱式注意力特征圖。作為MDTA的另一個(gè)重要成分，在計(jì)算特征協(xié)方差生成全局特征圖之前，我們引入了深度卷積以突出局部上下文。

對于輸入 , MDTA的處理流程如下:

• 首先, 生成 采用 深度卷積實(shí)現(xiàn), 注：bias-free);
• 然后, 對query與key進(jìn)行reshape并進(jìn)行點(diǎn)乘生成轉(zhuǎn)置注意力圖

MDTA的計(jì)算公式可以表示為：

注：是一個(gè)可學(xué)習(xí)縮放參數(shù)，用于控制點(diǎn)積的幅值。

Gated-Dconv Feed-forward Network

為進(jìn)行特征變換, 常規(guī)FFN(由兩個(gè) 卷積構(gòu)成)對每個(gè)像素進(jìn)行獨(dú)立處理。本文則對齊進(jìn)行了兩個(gè)改動(dòng)以 提升表達(dá)能力(見上圖)：

• gating機(jī)制：它表現(xiàn)為兩個(gè)并行分支的元素乘操作;
• 深度卷積：它表現(xiàn)為 深度卷積以學(xué)習(xí)局部結(jié)構(gòu)信息。

總而言之, GDFN可以控制信息流動(dòng), 進(jìn)而使得每層聚焦于不同的細(xì)節(jié)信息。由于GDFN具有比FFN更多的 計(jì)算量, 我們降低擴(kuò)展因子火以控制參數(shù)量與計(jì)算量。

Progressive Learning

基于CNN的復(fù)原模型通過采用固定尺寸圖像塊進(jìn)行訓(xùn)練。然而，Transformer模型在較小塊上訓(xùn)練可能無法進(jìn)行全局統(tǒng)計(jì)信息編碼，進(jìn)而導(dǎo)致全分辨率測試時(shí)的次優(yōu)性能。

針對該問題，我們提出了Progressive Learning機(jī)制：在訓(xùn)練的初期，模型在較小圖像塊上進(jìn)行訓(xùn)練；在訓(xùn)練的后期，模型采用更大圖像塊進(jìn)行訓(xùn)練 。由于更大的圖像塊會(huì)導(dǎo)致更長的計(jì)算耗時(shí)，我們隨圖像塊提升降低batch參數(shù)以保持與固定塊訓(xùn)練相當(dāng)?shù)暮臅r(shí)。

通過混合尺寸圖像塊訓(xùn)練的模型具有更優(yōu)的性能。Progressive學(xué)習(xí)策略具有類似Curriculum學(xué)習(xí)策略相似的行為。

Experiments

我們在不同的任務(wù)上進(jìn)行了所提方案的性能驗(yàn)證，包含圖像去雨、圖像去運(yùn)動(dòng)模糊、圖像去散焦模糊、圖像降噪。

Image Deraining

上圖&表給出了所提方案在去雨任務(wù)上的性能與效果對比，可以看到：

• 相比此前最佳SPAIR，Restormer在所有數(shù)據(jù)集上取得了平均1.05dB指標(biāo)提升 ；
• 在Rain100L數(shù)據(jù)集上，性能增益甚至高達(dá)2.06dB ；
• Restormer可以生成更好的無雨圖像，且可以有效的保持結(jié)構(gòu)內(nèi)容。

Motion Deblurring

上表給出了所提方案在不同去模糊數(shù)據(jù)集上的性能對比，可以看到：

• 相比MIMO-UNet+，所提Restormer可以取得了平均0.47dB指標(biāo)提升；
• 相比MPRNet，所提方案Restormer可以取得平均0.26dB指標(biāo)提升；Restormer的FLOPs僅為MPRNet的81%；
• 相比IPT，所提方案Restormer取得了0.4dB指標(biāo)提升，同時(shí)具有更少的參數(shù)量(4.4x)、更快的推理速度(29x)；
• 可視化效果見下圖，很明顯：Restormer重建結(jié)果更銳利清晰。

Defocus Deblurring

上圖&表給出了去散焦模糊任務(wù)上的性能對比，從中可以看到：

• 無論是單幀還是雙攝圖像，所提方案大幅均優(yōu)于其他方案；
• 在組合場景方面，相比此前最佳IFAN，所提方案取得了約0.6dB指標(biāo)提升；相比Uformer，所提方案取得了1.01dB指標(biāo)提升；
• 所提方案可以有效移除空間可變的散焦模糊。

Image Denoising

上表&圖給出了不同圖像降噪任務(wù)上的性能與效果對比，從中可以看到：

• Gaussian Denoising：所提方案在兩種實(shí)驗(yàn)配置下均取得了SOTA性能。對于極具挑戰(zhàn)性的50噪聲水平的Urban100數(shù)據(jù)，Restormer取得了比DRUNet高0.37dB指標(biāo)，比SwinIR高0.31dB。此外，相比SwinIR，Restormer計(jì)算量更少，速度更快。
• Real Denoising：所提方案是僅有的指標(biāo)超過40dB的方案。相比此前最佳MIRNet與Uformer，所提方案分貝取得了0.3dB與0.25dB指標(biāo)提升。
• Visual：從視覺效果上來看，所提方案可以重建更清晰的圖像，同時(shí)具有更細(xì)粒度的紋理。

Ablation Study

上表給出了關(guān)于模塊的消融實(shí)驗(yàn)，可以看到：

• 相比基線模型，MDTA可以帶來0.32dB指標(biāo)提升
• GDFN可以在MTA基礎(chǔ)上取得0.26dB指標(biāo)提升；
• GDFN與MDTA的組合可以取得了0.51dB指標(biāo)提升。

從上表可以看到：

• 在提煉階段添加Transformer可以進(jìn)一步提升模型性能；
• Progressive學(xué)習(xí)機(jī)制可以取得更佳的指標(biāo)，提升約0.07dB；
• 深而窄的模型比寬而淺的模型質(zhì)保更高。

公眾號后臺(tái)回復(fù)“數(shù)據(jù)集”獲取50+深度學(xué)習(xí)數(shù)據(jù)集下載～