Restormer: 刷新多个low-level任务指标，性能超SwinIR与Uformer等Transformer复原方案

AIWalker

发布于 2021-12-02 05:31:48

2.6K0

文章被收录于专栏：AIWalkerAIWalker

CV干货，第一时间送达

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

code](https://github.com/swz30/Restormer

本文是MPRNet与MIRNet的作者在图像复原领域的又一力作，也是Transformer技术在low-level领域的又一个SOTA。针对Transformer在高分辨率图像复原中存在的难点，提出了两种MDTA与GDFN两种改进，极大程度上缓解了计算量与GPU缓存占用问题。所提方案刷新了多个图像复原任务的SOTA性能。

Abstract

通过MHSA与FFN进行改进，本文提出一种高效Transformer，它可以捕获长距离像素相关性，同时可适用于大尺寸图像。所提方案Restormer(Restoration Transformer)在多个图像复原任务上取得了SOTA性能，包含图像去雨、图像去运动模糊、图像去散焦模糊以及图像降噪(包含合成与真实噪声)，可参见下图。

本文主要贡献包含以下几点：

提出了一种编解码Transformer用于高分辨率图像上多尺度local-global表达学习，且无需进行局部窗口拆分；
提出一种MDTA(Multi-Dconv head Transposed Attention)模块，它有助于进行局部与非局部相关像素聚合，可以高效的进行高分辨率图像处理；
提出一种GDFN(Gated-Dconv Feed-forward Network)模块，它可以执行可控特征变换，即抑制低信息特征，仅保留有用信息。

Method

本文旨在设计一种高效Transformer模型，它可以处理复原任务中的高分辨率图像。为缓解计算瓶颈，我们对MHSA进行了关键性改进并引入多尺度分层模块，见下图。

Overall Pipeline 给定退化图像I \in R^{H \times W \times 3} ，Restormer的处理流程如下：

首先，采用卷积得到底层特征F_0 \in R^{H \times W \times 3} ；
然后，这些底层特征将经由4级对称编解码器处理变换为深层特征F_d \in R^{H \times W \times 2C}

（每级编解码器包含多个Transformer模块，编码器分层降低空间分辨率，同时扩展通道，解码器则降低通道数，提升空间分辨率），注：为辅助重建过程，编解特征将与解码特征通过跳过连接进行concat并后接1\times 1

卷积降维;

其次，深层特征将经由提炼模块在进一步丰富特征多样性；
最后，通过卷积层处理生成残差图像R \in R^{H \times W \times 3} ，与输入图像相加即可得到复原图像\hat{I} = I + R 。

Multi-Dconv Head Transposed Attention

Transformer的主要计算负担来自自注意力层。在常规自注意力中，key-query的点乘操作会随分辨率以\mathcal{O}(W^2H^2)

增长。因此，基本不可能将自注意力层用于图像复原任务(图像复原的输入分辨率一般都比较大)。

为解决该问题，我们提出了MDTA(见上图)，它具有线性复杂度，其关键成分在于：在通道维度(而非空间维度)执行自注意力计算跨通道的交叉协方差以生成关于全局上下文的隐式注意力特征图。

作为MDTA的另一个重要成分，在计算特征协方差生成全局特征图之前，我们引入了深度卷积以突出局部上下文。

对于输入Y \in R^{\hat{H} \times \hat{W} \times \hat{C}} ，MDTA的处理流程如下：

首先，生成Q、K与V(采用1\times 1 + 3\times 3 深度卷积实现，注：bias-free)；
然后，对query与key进行reshape并进行点乘生成转置注意力图A \in R^{\hat{C} \times \hat{C}}

MDTA的计算公式可以表示为：

\hat{X} = W_p Attention(\hat{Q}, \hat{K}, \hat{V}) + X

Attention(\hat{Q}, \hat{K}, \hat{V}) = \hat{V} \cdot Softmax(\hat{K} \cdot \hat{Q}/\alpha)

注：\alpha 是一个可学习缩放参数，用于控制点积的幅值。

Gated-Dconv Feed-forward Network

为进行特征变换，常规FFN(由两个1\times 1 卷积构成)对每个像素进行独立处理。本文则对齐进行了两个改动以提升表达能力(见上图)：

gating机制：它表现为两个并行分支的元素乘操作；
深度卷积：它表现为3\times 3 深度卷积以学习局部结构信息。

GDFN的处理过程可以表示为如下公式：

\hat{X} = W_p^0 Gating(X) + X

Gating(X) = \phi(W_d^1 W_p^1(LN(X))) \odot W_d^2 W_p^2(LN(X))

总而言之，GDFN可以控制信息流动，进而使得每层聚焦于不同的细节信息。由于GDFN具有比FFN更多的计算量，我们降低扩展因子\gamma 以控制参数量与计算量。

Progressive Learning

基于CNN的复原模型通过采用固定尺寸图像块进行训练。然而，Transformer模型在较小块上训练可能无法进行全局统计信息编码，进而导致全分辨率测试时的次优性能。

针对该问题，我们提出了Progressive Learning机制：在训练的初期，模型在较小图像块上进行训练；在训练的后期，模型采用更大图像块进行训练 。由于更大的图像块会导致更长的计算耗时，我们随图像块提升降低batch参数以保持与固定块训练相当的耗时。

通过混合尺寸图像块训练的模型具有更优的性能。Progressive学习策略具有类似Curriculum学习策略相似的行为。

Experiments

我们在不同的任务上进行了所提方案的性能验证，包含图像去雨、图像去运动模糊、图像去散焦模糊、图像降噪。

Image Deraining

上图&表给出了所提方案在去雨任务上的性能与效果对比，可以看到：

相比此前最佳SPAIR，Restormer在所有数据集上取得了平均1.05dB指标提升 ；
在Rain100L数据集上，性能增益甚至高达2.06dB ；
Restormer可以生成更好的无雨图像，且可以有效的保持结构内容。

Motion Deblurring

上表给出了所提方案在不同去模糊数据集上的性能对比，可以看到：

相比MIMO-UNet+，所提Restormer可以取得了平均0.47dB指标提升；
相比MPRNet，所提方案Restormer可以取得平均0.26dB指标提升；Restormer的FLOPs仅为MPRNet的81%；
相比IPT，所提方案Restormer取得了0.4dB指标提升，同时具有更少的参数量(4.4x)、更快的推理速度(29x)；
可视化效果见下图，很明显：Restormer重建结果更锐利清晰。

Defocus Deblurring

上图&表给出了去散焦模糊任务上的性能对比，从中可以看到：

无论是单帧还是双摄图像，所提方案大幅均优于其他方案；
在组合场景方面，相比此前最佳IFAN，所提方案取得了约0.6dB指标提升；相比Uformer，所提方案取得了1.01dB指标提升；
所提方案可以有效移除空间可变的散焦模糊。

Image Denoising

上表&图给出了不同图像降噪任务上的性能与效果对比，从中可以看到：

Gaussian Denoising：所提方案在两种实验配置下均取得了SOTA性能。对于极具挑战性的50噪声水平的Urban100数据，Restormer取得了比DRUNet高0.37dB指标，比SwinIR高0.31dB。此外，相比SwinIR，Restormer计算量更少，速度更快。
Real Denoising：所提方案是仅有的指标超过40dB的方案。相比此前最佳MIRNet与Uformer，所提方案分贝取得了0.3dB与0.25dB指标提升。
Visual：从视觉效果上来看，所提方案可以重建更清晰的图像，同时具有更细粒度的纹理。