《Restormer：高效Transformer架构用于高分辨率图像恢复》学习笔记

paper：2111.09881

GitHub：swz30/Restormer： [CVPR 2022--Oral] Restormer：高分辨率图像修复的高效转换器。SOTA 用于运动去模糊、图像去模糊、去噪（高斯/真实数据）和去焦去模糊。

复现：Restormer 代码复现_restormer复现-CSDN博客

摘要

1、介绍

2、相关工作

2.1 图像恢复

2.2 视觉Transformer

3、方法

3.1 多维卷积头转置注意力（MDTA）

3.2 门控深度卷积前馈网络（GDFN）

3.3 渐进式学习

4、实验

4.1 图像去雨结果

4.2 单图像运动去模糊结果

4.3 焦外模糊去模糊结果

4.4 图像去噪结果

4.5 消融实验

摘要

由于卷积神经网络（CNN）能够从大规模数据中学习出具有良好泛化性的图像先验，因此这些模型已经被广泛应用于图像恢复及相关任务。最近，另一类神经网络架构——Transformer，在自然语言处理和高层次视觉任务上展现出了显著的性能提升。虽然 Transformer 模型在缓解CNN的缺点（例如有限的感受野和对输入内容的适应性差）方面有优势，但其计算复杂度随着空间分辨率的增加呈平方级增长，因此在涉及高分辨率图像的大多数图像恢复任务中应用起来非常困难。在本研究中，我们通过在构建模块（多头注意力和前馈网络）中进行若干关键设计，提出了一种高效的Transformer模型，使其能够捕捉长程像素交互，同时仍然能够应用于大尺寸图像。我们提出的模型命名为Restoration Transformer（Restormer），在多个图像恢复任务上取得了最先进的结果，包括图像去雨、单幅图像运动去模糊、焦外去模糊（单幅图像和双像素数据）、以及图像去噪（高斯灰度/彩色去噪和真实图像去噪）。

1、介绍

图像恢复是通过去除退化（如噪声、模糊、雨滴等）来重建高质量图像的任务。由于其不适定的特性，这个问题非常具有挑战性，通常需要强大的图像先验来实现有效恢复。由于卷积神经网络（CNN）能够从大规模数据中学习到具有良好泛化性的先验，因此它们已经成为相较于传统恢复方法的首选。

CNN的基本操作是“卷积”，它提供了局部连接性和位移等变性。虽然这些特性为CNN带来了效率和泛化能力，但也引发了两个主要问题：

（a）卷积操作具有有限的感受野，从而无法建模长程像素依赖关系；

（b）卷积滤波器在推理时具有静态权重，因此无法灵活地适应输入内容。

为了应对上述问题，一种更强大且动态的替代方案是自注意力（SA）机制[17, 77, 79, 95]，它通过对所有其他位置进行加权求和来计算给定像素的响应。（背景）

自注意力是 Transformer 模型的核心组件[34, 77]，并且有独特的实现方式，即多头SA，它优化了并行化和有效的表示学习。Transformer在自然语言任务[10, 19, 49, 62]和高层次视觉问题[11, 17, 76, 78]中展示了最先进的性能。虽然SA在捕捉长程像素交互方面非常有效，但其复杂度随着空间分辨率的增加呈平方级增长，因此使其在处理高分辨率图像（图像恢复中常见的情况）时变得不可行。

最近，少数工作尝试将 Transformer 模型定制化，以应用于图像恢复任务[13, 44, 80]。为了减少计算负担，这些方法要么在每个像素周围应用8x8大小的SA小空间窗口[44, 80]，要么将输入图像分割成48x48大小的不重叠小块，并独立计算每个小块的SA[13]。然而，限制SA的空间范围与捕捉真实的长程像素关系的目标相矛盾，尤其是在高分辨率图像上。（Transformer 的问题）

在本文中，我们提出了一种高效的Transformer模型，用于图像恢复，能够建模全局连接性，并且仍然适用于大尺寸图像。具体来说，我们引入了一种多D卷积头“转置”注意力（MDTA）模块（第3.1节），代替了传统的多头SA[77]，其复杂度为线性。MDTA跨越特征维度应用SA，而不是空间维度；即它不显式地建模像素间的配对交互，而是计算跨特征通道的协方差，以从（键和值投影）输入特征中获得注意力图。MDTA模块的一个重要特征是在计算特征协方差之前进行局部上下文混合。这是通过使用1x1卷积进行像素级别的跨通道上下文聚合，以及使用高效的深度卷积进行通道级局部上下文聚合来实现的。这种策略带来了两个关键优势。

（1）它强调空间局部上下文，并在我们的管道中引入了卷积操作的互补强度。

（2）它确保了像素间的全局关系在计算基于协方差的注意力图时被隐式建模。

前馈网络（FN）是Transformer模型的另一个构建模块[77]，它由两个全连接层组成，并且中间有非线性激活。在这项工作中，我们通过引入门控机制[16]重新构造了常规FN[77]的第一个线性变换层，以提高信息流动性。这个门控层被设计为两个线性投影层的逐元素积，其中一个通过GELU非线性激活[27]。我们的门控D卷积FN（GDFN）（第3.2节）也基于类似MDTA模块的局部内容混合，均衡地强调空间上下文。GDFN中的门控机制控制哪些互补特征应该传递下去，并允许网络层次中的后续层专注于更精细的图像属性，从而产生高质量的输出。

门控机制的核心是通过非线性函数（如sigmoid）生成权重矩阵，对输入信息进行选择性过滤或加权融合。

除了上述架构创新外，我们还展示了 Restormer 的渐进式学习策略的有效性（第3.3节）。在这个过程中，网络在早期 epochs 中对小的图像块和大批量进行训练，在后期epochs中对逐渐增大的图像块和小批量进行训练。这种训练策略帮助Restormer从大图像中学习上下文，并在测试时提供质量性能提升。我们进行了全面的实验，并展示了Restormer在16个基准数据集上的最先进性能，涵盖了多个图像恢复任务，包括图像去雨、单幅图像运动去模糊、焦外去模糊（单幅图像和双像素数据）以及图像去噪（在合成和真实数据上）。此外，我们提供了大量消融实验，证明了架构设计和实验选择的有效性。

主要贡献总结：

我们提出了Restormer，一个用于高分辨率图像的编码器-解码器Transformer，通过多尺度局部-全局表示学习来处理图像，且无需将图像分割为局部窗口，从而充分利用远程图像上下文。
我们提出了一个多D卷积头转置注意力（MDTA）模块，能够聚合局部和非局部像素交互，并且足够高效，能够处理高分辨率图像。
我们提出了一种新的门控D卷积前馈网络（GDFN），能够进行受控的特征转换，即抑制不太有信息的特征，并允许只有有用的信息向网络层次继续流动。

2、相关工作

2.1 图像恢复

近年来，基于数据驱动的卷积神经网络（CNN）架构[7, 18, 92, 93, 105, 107]已被证明优于传统的恢复方法[26, 36, 53, 75]。在卷积设计中，基于编码器-解码器的U-Net架构[3, 14, 39, 80, 90, 93, 99]由于其分层的多尺度表示，同时保持计算效率，一直是恢复任务中的主流设计。类似地，基于跳跃连接的方法[24, 48, 92, 106]也因专注于学习残差信号而被证明在恢复任务中有效。空间和通道注意力模块也被引入以选择性地关注相关信息[43, 92, 93]。我们建议读者参考NTIRE挑战报告[2, 5, 30, 57]和最近的文献综述[8, 42, 73]，这些文献总结了图像恢复中的主要设计选择。

2.2 视觉Transformer

Transformer模型最初是为自然语言任务中的序列处理而开发的[77]。它已被应用于许多视觉任务，如图像识别[17, 76, 88]、分割[78, 83, 108]、目标检测[11, 50, 109]等。Vision Transformer（ViT）[17, 76]将图像分解为一系列 patch（局部窗口），并学习这些 patch 之间的相互关系。这些模型的显著特点是能够强有力地学习图像 patch 序列之间的长程依赖关系，并能适应给定的输入内容[34]。由于这些特性，Transformer模型也被应用于低级视觉问题，如超分辨率[44, 85]、图像着色[37]、去噪[13, 80]、去雨[80]等。然而，Transformer中自注意力（SA）的计算复杂度随着图像补丁数目的增加呈平方级增长，这使得其在高分辨率图像中的应用变得不可行。因此，在需要生成高分辨率输出的低级图像处理应用中，最近的方法通常采用不同的策略来降低计算复杂度。一种可能的解决方案是使用Swin Transformer设计[44]，在局部图像区域内应用自注意力[44, 80]。然而，这种设计选择将上下文聚合限制在局部邻域内，违背了使用自注意力而非卷积的主要动机，因此不太适合图像恢复任务。相比之下，我们提出了一种Transformer模型，能够在保持计算效率的同时，学习长程依赖关系。

3、方法

我们的主要目标是开发一个高效的 Transformer 模型，能够处理高分辨率图像用于恢复任务。为了缓解计算瓶颈，我们在多头自注意力层和多尺度分层模块中引入了关键设计，使得这些模块的计算需求低于单尺度网络[44]。我们首先展示了Restormer架构的整体流程（见图2）。然后，我们描述了所提出的Transformer模块的核心组件：

（a）多维卷积头转置注意力（MDTA）

（b）门控卷积前馈网络（GDFN）

最后，我们提供了关于渐进式训练策略的细节，以有效学习图像统计信息。

整体流程：给定一张降质图像 $I \in \mathbb{R}^{H \times W \times 3}$ ，Restormer 首先应用卷积操作以获得低级特征嵌入 $F_0 \in \mathbb{R}^{H \times W \times C}$ ，其中 $H$ 和 $W$ 是空间维度， $C$ 是通道数。接下来，这些浅层特征 $F_0$ 通过一个4层对称的编码器-解码器，并被转换为深层特征 $F_d \in \mathbb{R}^{H \times W \times 2C}$ 。每一层的编码器-解码器包含多个Transformer模块，并且从顶部到底部的模块数量逐渐增加，以保持计算效率。编码器从高分辨率输入开始，逐步减少空间尺寸，同时增加通道容量。解码器接收低分辨率的潜在特征 $F_l \in \mathbb{R}^{\frac{H}{8} \times \frac{W}{8} \times C}$ 作为输入，并逐步恢复高分辨率的表示。为了特征下采样和上采样，我们分别应用像素反打乱（pixel-unshuffle）和像素打乱（pixel-shuffle）操作[69]。为了辅助恢复过程，编码器的特征通过跳跃连接与解码器的特征进行拼接[66]。拼接操作后，所有层（除了顶部层）都会通过 1×1 卷积来减少通道数（减半）。在第1层，我们让Transformer模块将编码器的低级图像特征与解码器的高级特征进行聚合，这有助于保持恢复图像中的细节结构和纹理。接下来，深层特征 $F_d$ 在细化阶段进一步丰富，并在高空间分辨率下进行处理。这些设计选择带来了质量上的提升，如我们在实验部分（第4节）中所见。最后，应用卷积层到细化后的特征上，生成残差图像 $R \in \mathbb{R}^{H \times W \times 3}$ ，并将降质图像加到残差图像上，得到恢复后的图像： $\hat{I} = I + R$ 。

3.1 多维卷积头转置注意力（MDTA）

Transformer中的主要计算开销来自于自注意力层。在传统的自注意力（SA）[17, 77]中，键-查询点积交互的时间和内存复杂度随着输入的空间分辨率的增加而成平方增长，即对于W×H像素的图像，复杂度为 $O(W^2 H^2)$ 。因此，在大多数涉及高分辨率图像的图像恢复任务中，应用自注意力是不可行的。为了缓解这一问题，我们提出了MDTA，如图2(a)所示，其具有线性复杂度。其关键要素是应用自注意力跨通道而不是跨空间维度，即计算跨通道的协方差以生成编码全局上下文的注意力图。

在MDTA的另一个重要组成部分中，我们引入了深度卷积（depth-wise convolutions）来在计算特征协方差之前强调局部上下文，从而产生全局注意力图。

从一个层归一化的张量 $Y \in \mathbb{R}^{H \times W \times C}$ 开始，MDTA首先生成查询（Q）、键（K）和值（V）投影，这些投影已通过局部上下文丰富。具体来说，我们通过应用 1×1 卷积聚合像素级的跨通道上下文，随后通过 3×3 深度卷积来编码通道级的空间上下文，得到：

$Q = W^Q_{d}W^Q_{p}Y, \quad K = W^K_{d}W^K_{p}Y, \quad V = W^V_{d}W^V_{p}Y$

其中， $W^Q_{p},W^K_{p},W^V_{p}$ 是 1×1 卷积， $W^Q_{d},W^K_{d},W^V_{d}$ 是 3×3深度卷积。接下来，我们将查询和键投影进行重塑，以使它们的点积交互生成一个转置的注意力图 A ，其大小为 $R ^ {\hat{C}\times \hat{C}}$ ，而不是传统的巨大注意力图 $R ^{ \hat{H}\hat{W} \times \hat{H}\hat{W}}$ [17, 77]。整体MDTA过程定义为：

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

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

其中 $X$ 和 $\hat{X}$ 分别是输入和输出特征图， $\hat{Q}\in \mathbb{R}^{\hat{H} \hat{W} \times \hat{C}},\hat{K}\in \mathbb{R}^{\hat{C}\times\hat{H} \hat{W} },\hat{V}\in \mathbb{R}^{\hat{H} \hat{W} \times \hat{C}}$ 是在重塑张量之后从原始大小 $R ^{ \hat{H} \times \hat{W} \times \hat{C}}$ 获得的矩阵。这里， $\alpha$ 是可学习的缩放参数，用来控制 K 和 Q 的点积大小，之后应用softmax函数。与传统的多头自注意力（multi-head SA）[17]类似，我们将通道数分成多个“头”，并并行学习独立的注意力图。

3.2 门控深度卷积前馈网络（GDFN）

为了变换特征，常规的前馈网络（FN）[17, 77]是针对每个像素位置独立且相同地操作的。它使用两个 1×1 卷积，一个用于扩展特征通道（通常扩展倍数为4），第二个则将通道数减少回原始输入维度。在隐藏层中应用非线性函数。在本工作中，我们提出了对FN的两项基本修改，以改善表示学习：

1）引入门控机制，

2）引入深度卷积。

我们的GDFN架构如图2(b)所示。

门控机制被公式化为两个并行线性变换路径的逐元素乘积，其中一个路径通过GELU非线性激活[27]。如同MDTA中一样，我们还在GDFN中使用了深度卷积来编码来自空间相邻像素位置的信息，这对学习局部图像结构以有效恢复非常有用。

给定输入张量 $X \in R ^{ \hat{H} \times \hat{W} \times \hat{C}}$ ，GDFN的公式为：

$\hat{X} = W^0_{p} \cdot Gating(X) + X$

$Gating(X) = \phi (W^1_{d}W^1_{p}(\text{LN}(X))) \bigodot (W^2_{d}W^2_{p}(\text{LN}(X)))$

其中， $\bigodot$ 表示逐元素乘法， $\phi$ 表示GELU非线性激活函数，LN 表示层归一化[9]。总体来说，GDFN控制着信息流动，通过我们管道中的不同层级，从而允许每个层专注于对其他层互补的细节内容。也就是说，GDFN与MDTA的作用不同，MDTA侧重于丰富特征并增加上下文信息，而GDFN则主要处理信息流动。

由于我们提出的GDFN相比常规FN[17]执行了更多的操作，因此我们减少了扩展比率，以使其具有类似的参数数量和计算负担。

3.3 渐进式学习

基于CNN的恢复模型通常在固定大小的图像补 patch 上进行训练。然而，在小的裁剪补丁上训练Transformer模型可能无法编码全局图像统计信息，从而在测试时对全分辨率图像的表现较差。为此，我们采用渐进式学习策略，在早期阶段使用较小的图像 patch 进行训练，并在后期阶段逐渐使用更大的图像 patch 。通过渐进式学习训练的模型在测试时表现更好，因为图像的分辨率可以不同（这在图像恢复中很常见）。渐进式学习策略类似于课程学习（curriculum learning）过程，网络从简单的任务开始，逐渐转向学习更复杂的任务（例如，保留精细的图像结构/纹理）。由于使用大 patch 训练会导致训练时间增加，我们在 patch 大小增大的同时减少了批量大小，以保持每个优化步骤的时间与固定 patch 训练的相似。

4、实验

我们在基准数据集和实验设置上对提出的Restormer进行了评估，涵盖四种图像处理任务：

(a) 图像去雨

(b) 单图像运动去模糊

(c) 焦外模糊去模糊（包括单图像和双像素数据）

(d) 图像去噪（包括合成和真实数据）

更多关于数据集、训练协议和额外视觉结果的细节可以在补充材料中找到。在表格中，评估方法的最佳和第二最佳质量得分被高亮并下划线显示。

实现细节

我们为不同的图像恢复任务训练了独立的模型。在所有实验中，除非另有说明，否则我们使用以下训练参数。我们的Restormer采用了4层编码器-解码器结构。从第1层到第4层，Transformer块的数量分别为[4,6,6,8]，MDTA中的注意力头数为[12,24,48]，通道数为[48,96,192,384]。细化阶段包含4个块。GDFN中的通道扩展因子为2.66。我们使用AdamW优化器（1=0.9，2=0.999，权重衰减为1e-4）和L1损失训练了300K次迭代，初始学习率为3e-4，并通过余弦退火逐渐降低到1e-6。对于渐进学习，我们从训练大小为128x128的图像开始，批大小为64。然后，图像的训练大小和批大小逐渐调整为[(1602,40)，(1922,32)，(2562,16)，(3202,8)，(3842,8)]，对应的迭代次数分别为[92K,156K,204K,240K,276K]。在数据增强方面，我们使用了水平和垂直翻转。

4.1 图像去雨结果

我们使用YCbCr颜色空间中的Y通道计算PSNR/SSIM得分，类似于现有方法[32,61,93]。表1显示，我们的Restormer在所有五个数据集上均表现出一致且显著的性能提升。与最近的最佳方法SPAIR[61]相比，Restormer在所有数据集上的平均PSNR提升了10.5 dB。在单个数据集上，增益可以达到20.6 dB，例如Rain100L。图3展示了一个具有挑战性的视觉示例。我们的Restormer有效去除了雨滴，同时保留了结构信息。

4.2 单图像运动去模糊结果

我们在合成数据集（GoPro [56]、HIDE [67]）和真实数据集（RealBlur-R[65]、RealBlur-J[65]）上评估了去模糊方法。表2显示，Restormer在所有四个基准数据集上均优于其他方法。在所有数据集上，Restormer相比最近的MIMO-UNet+[14]方法平均提升了0.47 dB，相比之前的最佳方法MPRNet [93]提升了0.26 dB。与MPRNet [93]相比，Restormer的FLOPs减少了81%（见图1）。此外，我们的方法相较于Transformer模型IPT[13]，提高了0.4 dB，而参数减少了44%，速度提高了29%。值得注意的是，我们的Restormer仅在GoPro[56]数据集上进行了训练，但它在其他数据集上也展示了强大的泛化能力，设定了新的SOTA。图4显示，我们的方法生成的图像比其他算法的结果更清晰，且视觉上更接近真实图像。

4.3 焦外模糊去模糊结果

表3展示了常规焦外模糊去模糊方法（EBDB[33]和JNB[68]）以及基于学习的方法在DPDD数据集[3]上的图像保真度得分。我们的Restormer在单图像和双像素焦外模糊去模糊任务上显著超越了所有的最先进方案，尤其在组合场景类别中，Restormer比之前的最佳方法IFAN[41]提高了0.6 dB。与Transformer模型Uformer[80]相比，我们的方法提供了10.1 dB的PSNR增益。图5表明，我们的方法在去除空间变化的焦外模糊方面比其他方法更有效。

4.4 图像去噪结果

我们在合成基准数据集（Set12[101]、BSD68[52]、Urban100[29]、Kodak24[20]和McMaster[104]）以及真实数据集（SIDD[1]和DND[60]）上进行了去噪实验，遵循[54,93,99]的流程，我们使用无偏Restormer进行去噪。

高斯去噪

表4和表5展示了不同方法在多个基准数据集上的PSNR得分，分别为灰度图像和彩色图像的去噪。在测试时，我们考虑了噪声级别15、25和50。评估方法分为两类实验设置： (1) 学习一个模型处理各种噪声级别，(2) 为每个噪声级别学习一个独立模型。我们的Restormer在这两种实验设置下，在不同数据集和噪声级别下均取得了SOTA的性能。特别是在高分辨率Urban100数据集[29]上，对于挑战性的噪声级别50，Restormer比之前的最佳CNN方法DRUNet [99]提高了0.37 dB，比最近的基于Transformer的网络SwinIR[44]提高了0.31 dB，见表4。类似的性能提升也可以在表5中观察到，针对高斯彩色去噪。值得一提的是，DRUNet[99]需要噪声级别图作为额外输入，而我们的方法只使用噪声图像。此外，相比SwinIR[44]，Restormer的FLOPs减少了31.4%，运行速度提升了13%。图6展示了不同方法的去噪结果，灰度去噪（上排）和彩色去噪（中排）。我们的Restormer恢复了清晰的图像。

真实图像去噪

表6显示，我们的方法是唯一一个在两个数据集上都超过40 dB PSNR的方法。值得注意的是，在SIDD数据集上，我们的Restormer相比之前的最佳CNN方法MIRNet[92]和Transformer模型Uformer[80]，分别取得了0.3 dB和0.25 dB的PSNR增益。图6（底排）显示，我们的Restormer在不妥协细节纹理的情况下生成了清晰的图像。

4.5 消融实验

在消融实验中，我们训练了一个高斯彩色去噪模型，图像块大小为128x128，训练迭代次数为100K。测试在Urban100[29]数据集上进行，分析集中于具有挑战性的噪声级别50。FLOPs和推理时间计算是在256x256图像大小上进行的。表7到表10展示了我们的贡献在质量性能上的提升。接下来，我们逐一描述每个组件的影响。

多头注意力的改进

表7c展示了我们的MDTA相较于基线方法（表7a）提供了0.32 dB的显著增益。此外，通过深度卷积引入局部性到MDTA，使其更加稳健，去除深度卷积后，PSNR出现了下降（见表7b）。

前馈网络（FN）的改进

表7d显示，在前馈网络（FN）中引入门控机制以控制信息流，提供了0.12 dB的增益，超越了传统的FN[77]。与多头注意力类似，在FN中引入局部机制也带来了性能上的优势（见表7e）。我们进一步通过引入门控深度卷积来增强FN的表现。我们的GDFN（见表7f）相比于标准的FN[77]，在噪声级别50下，取得了0.26 dB的PSNR增益。总体来说，我们的Transformer块贡献使得性能较基线方法提升了0.51 dB。

解码器设计（第1层）

为了在第1层将编码器特征与解码器特征聚合，我们在拼接操作后不使用1x1卷积（该操作会将通道数减少一半）。这种设计有助于保留来自编码器的细节信息，尤其是纹理细节，正如表8所示。这些结果进一步证明，在细化阶段添加Transformer块的有效性。

渐进学习的影响

表9显示，渐进学习相比固定补丁训练提供了更好的结果，且训练时间相似。

更深或更宽的Restormer？

表10展示了在相同参数/FLOPs预算下，深度-窄型模型的表现优于其宽浅型对比模型。然而，宽型模型因并行化优势而运行更快。在本文中，我们使用了深度-窄型的Restormer。