图像复原 --- Restormer

系列文章目录

---

文章名称：Restormer: Efficient Transformer for High-Resolution Image Restoration
文章地址：https://arxiv.org/abs/2111.09881
代码地址：https://github.com/swz30/restormer
发表时间：2022
引用领域：图像去雨、单图像运动去模糊、散焦去模糊（单图像和双像素数据）、图像去噪（高斯灰度/彩色图像/真实图像）
主要模块：MDTA、GDFN

---

---

摘要

由于卷积神经网络 (CNN) 在从大规模数据中学习可泛化图像先验方面表现良好，因此这些模型已广泛应用于图像恢复和相关任务。最近，另一类神经架构 Transformers 在自然语言和高级视觉任务上显示出显着的性能提升。虽然 Transformer 模型减轻了 CNN 的缺点（即有限的感受野和对输入内容的不适应性），但其计算复杂度随空间分辨率呈二次方增长，因此无法应用于大多数图像恢复涉及高分辨率图像的任务。在这项工作中，我们通过在构建块（多头注意力和前馈网络）中进行几个关键设计，提出了一个有效的 Transformer 模型，使其能够捕获远程像素交互，同时仍然适用于大图像。我们的模型名为 Restoration Transformer (Restormer)，在多个图像恢复任务上取得了最先进的结果，包括图像去雨、单图像运动去模糊、散焦去模糊（单图像和双像素数据）、和图像去噪（高斯灰度/颜色去噪和真实图像去噪）。源代码和预训练模型可在 https://github.com/swz30/Restormer 获得。

Introduction

图像恢复是通过去除退化输入中的退化（如噪声、模糊、雨滴）来重建高质量图像的任务。 由于图像的不适定性，通常需要较强的图像先验信息才能有效恢复，这是一个极具挑战性的问题。 由于卷积神经网络（CNNs）在从大规模数据中学习可推广的先验知识方面表现出色，因此与传统的恢复方法相比，卷积神经网络是一种更好的选择。
CNN:
CNNS中的基本操作是“卷积”，它提供局部conectivity和translation等效性。 尽管这些特性为CNNS带来了效率和泛化，但它们也带来了两个主要问题 。（a）卷积操作的感受野有限，因此无法对远程像素相关性进行建模。 (b) 卷积滤波器在推理时具有静态权重，因此不能灵活地适应输入内容。 为了解决上述缺点，一个更强大和动态的替代方案是自我注意 (SA) 机制 Trasformer、Non-local neural networks、Self-attention generative adversarial networks，它通过所有其他位置的加权和计算给定像素的响应 .
SA:
自注意是Transformer模型中的一个核心组成部分，但有一个独特的实现，即为并行化和有效的表示学习而优化的多头SA。 Transformer在自然语言任务和高级视觉问题上表现出了最先进的性能。 尽管SA在捕捉远距离像素相互作用方面非常有效，但其复杂度随空间分辨率呈二次增长，因此不适用于高分辨率图像（图像恢复中经常出现的情况）。 最近，很少有人为图像恢复任务量身定制 Transformerpretrained-IPT、SwinIR、Uformer。 为了减少计算量，这些方法要么将SA应用于每个像素周围大小为8×8的小空间窗口，或者是将输入图像分成大小为 48×48 的非重叠块，并在每个块上独立计算 SA [13]。 然而，限制 SA 的空间范围与捕获真正的远程像素关系的目标是矛盾的，尤其是在高分辨率图像上。
MDTA:
在本文中，我们提出了一种高效的图像恢复 Transformer，它能够对全局连接进行建模，并且仍然适用于大图像。 具体来说，我们引入了一个多 Dconv 头“转置”注意力 (MDTA) 块（第 3.1 节）来代替具有线性复杂度的普通多头 SA 。 它跨特征维度而不是空间维度应用 SA，即，MDTA 不是显式建模成对像素交互，而是跨特征通道计算交叉协方差，以从（key和query投影的）输入特征中获取注意力图。 我们的 MDTA 块的一个重要特征是特征协方差计算之前的局部上下文混合。 这是通过使用 1×1 卷积的跨通道上下文的像素级聚合和使用有效的深度卷积的局部上下文的通道级聚合来实现的。 该策略提供了两个关键优势。 首先，它强调空间局部上下文，并在我们的pipeline中引入卷积运算的优势。 其次，它确保在计算基于协方差的注意力图时隐式建模像素之间的上下文的全局关系。
GDFN:
前馈网络 (FN) 是 Transformer 模型的另一个构建块，它由两个完全连接的层组成，其间具有非线性。 在这项工作中，我们用门控机制 Language modeling with gated convolutional networks 重新制定了常规 FN 的第一个线性变换层，以改善通过网络的信息流。 该门控层被设计为两个线性投影层的逐元素乘积，其中一个由 GELU 非线性激活 。 我们的门控 Dconv FN (GDFN)（第 3.2 节）也基于类似于 MDTA 模块的本地内容混合，以同样强调空间上下文。 GDFN 中的门控机制控制哪些互补特征应该向前流动，并允许网络层次结构中的后续层专门关注更精细的图像属性，从而产生高质量的输出。
除了上述架构新颖性之外，我们还展示了 Restormer 渐进式学习策略的有效性（第 3.3 节）。 在此过程中，网络在早期阶段接受小块和大批量训练，在后期阶段逐渐接受大图像块和小批量训练。 这种训练策略帮助 Restormer 从大图像中学习上下文，并随后在测试时提供质量性能改进。 我们进行了全面的实验，并在 16 个基准数据集上展示了我们的 Restormer 的最先进性能，用于多项图像恢复任务，包括图像去雨、单图像运动去模糊散焦去模糊（在单图像和双像素数据上）和图像去噪（在合成和真实数据上）； 参见图 1。此外，我们提供了广泛的消融以显示架构设计和实验选择的有效性。 这项工作的主要贡献总结如下：
• 我们提出了 Restormer，一种编码器-解码器 Transformer，用于对高分辨率图像进行多尺度局部-全局表示学习，而无需将它们分解到局部窗口中，从而利用远距离图像上下文。
• 我们提出了一个多Dconv 头部转置注意(MDTA) 模块，该模块能够聚合局部和非局部像素交互，并且足够高效以处理高分辨率图像。
• 一种新的门控 Dconv 前馈网络 (GDFN)，它执行受控的特征转换，即抑制信息量较少的特征，并只允许有用的信息进一步通过网络层次结构。

Background

Image Restoration.
近年来，数据驱动的 CNN 架构Real image denoising with
feature attention、Burst image restoration and
enhancement、Learning enriched features for real image restoration
and enhancement、MPRNet、RDN已被证明优于传统的恢复方法。 在卷积设计中，基于编码器-解码器的 UNet 架构DNN(散焦去模糊）、MIMO-UNet（去模糊）、DeblurGAN-v2（去模糊）、DANet、MPRNet由于其分层多尺度表示同时保持计算效率，主要用于恢复研究。 同样，由于特别关注学习残差信号，基于跳跃连接的方法已被证明对恢复有效DualResidualNetwork、MIRNet、RNAN。 空间和通道注意模块也被合并以选择性地关注相关信息 [43,92,93]。 我们建议读者参考 NTIRE 挑战报告 [2,5,30,57] 和最近的文献综述 resoulution、derain、denosing，它们总结了图像恢复的主要设计选择。
Vision Transformer
Transformer 模型最初是为自然语言任务中的序列处理而开发的。 它已适用于许多视觉任务，例如图像识别、分割、目标检测。 Vision Transformers 将图像分解为一系列补丁（局部窗口）并了解它们之间的相互关系。 这些模型的显着特征是学习图像块序列之间的远程依赖关系的强大能力以及对给定输入内容的适应性 Transformers in Vision: A Survey。 由于这些特性，Transformer 模型也被研究用于低级视觉问题，例如超分辨率、图像着色、去噪 pretrained-IPT、Uformer和去雨 。 然而，Transformers 中 SA 的计算复杂度会随着图像块的数量呈二次方增长，从而禁止其应用于高分辨率图像。
最近的方法通常采用不同的策略来降低复杂性。 一种潜在的补救措施是使用 Swin Transformer 设计 [44] 在局部图像区域 SwinIR 、Uformer中应用自注意力。 然而，这种设计选择限制了局部邻域内的上下文聚合，违背了使用自注意力而不是卷积的主要动机，因此不适合图像恢复任务。 相比之下，我们提出了一个 Transformer 模型，它可以在保持计算效率的同时学习远程依赖关系。

Method

我们的主要目标是开发一种高效的 Transformer 模型，可以处理用于修复任务的高分辨率图像。 为了缓解计算瓶颈，我们将关键设计引入多头 SA 层和多尺度分层模块，其计算要求低于单尺度网络 SwinIR。 我们首先展示我们的 Restormer 架构的整体pipeline（见图 2）。 然后我们描述了所提出的 Transformer 块的核心组件：(a) 多 Dconv 头部转置注意力 (MDTA) 和 (b) 门控 Dconv 前馈网络 (GDFN)。 最后，我们提供了有关有效学习图像统计的渐进式训练方案的详细信息。

Overall Pipeline.(整体网络架构）

给定退化图像 $\mathbf{I}\in {\mathbb{R}}^{H\times W\times 3}$，Restormer 首先应用卷积以获得低级特征嵌入 ${\mathbf{F}}_0\in {\mathbb{R}}^{H\times W\times C}$； 其中 H × W 表示空间维度，C 是通道数。
接下来，这些浅层特征 $F_0$ 通过一个 4 级对称编码器-解码器，转化为深层特征 ${\mathbf{F}}_{\mathbf{d}}\in {\mathbb{R}}^{H\times W\times 2C}$ 。
encoder-decoder的每一层都包含多个Transformer block，其中block的个数从上到下逐渐增加以保持效率。 从高分辨率输入开始，编码器分层缩小空间尺寸，同时扩展通道容量。 解码器将低分辨率潜在特征${\mathbf{F}}_l\in {\mathbb{R}}^{\frac{H}{8}\times \frac{W}{8}\times 8C}$作为输入，并逐步恢复高分辨率表示。
对于特征下采样和上采样，我们分别应用pixel-unshuffle和pixel-shuffle操作subpixel。 为了协助恢复过程，编码器特征通过skip-connection与解码器特征连接U-net。 连接操作之后是 1×1 卷积，以减少所有级别的通道（减半），除了顶部。
在第 1 级，我们让 Transformer 块将编码器的低级图像特征与解码器的高级特征进行聚合。 它有利于在恢复的图像中保留精细的结构和纹理细节。 接下来，深层特征 Fd 在以高空间分辨率运行的细化阶段进一步丰富。 正如我们将在实验部分（第 4 节）中看到的那样，这些设计选择会带来质量改进。
最后，将卷积层应用于细化特征以生成残差图像 $\mathbf{R}\in {\mathbb{R}}^{H\times W\times 3}$，向其添加退化图像以获得恢复图像：$\hat{\mathbf{I}}=\mathbf{I}+\mathbf{R}$。接下来，我们介绍 Transformer 模块的模块 .

Multi-Dconv Head Transposed Attention（MDTA)

Transformers中的主要开销来自自注意力层。在传统的SA中，key、query点积交互的时间和空间复杂度随着输入的空间分辨率呈二次方增长，即对W*H像素的图像，时间复杂度为$\mathcal{O}(W^2{H}^2)$,因此，将SA应用于大多数通常涉及高分辨率图像的恢复任务是不可行的。为了缓解这个问题，我们提出了具有线性复杂度的 MDTA，如图 2(a) 所示。 关键要素是跨通道而不是空间维度应用 SA，即计算跨通道的互协方差以生成隐式编码全局上下文的注意力图。 作为 MDTA 的另一个重要组成部分，我们引入了深度卷积以在计算特征协方差以生成全局注意力图之前强调局部上下文。
从层归一化张量$\mathbf{Y}\in {\mathbb{R}}^{\hat{H}\times \hat{W}\times \hat{C}}$ ，我们的 MDTA 首先生成查询 (Q)、键 (K) 和值 (V) 投影，并丰富了本地上下文。 它是通过应用 1×1 卷积来聚合像素级跨通道上下文，然后使用 3×3 深度卷积来编码通道级空间上下文来实现的，得到 $\mathbf{Q}=W_d^Q{W}_p^Q\mathbf{Y},\mathbf{K}=W_d^K{W}_p^K\mathbf{Y}\mathrm{and}\mathbf{V}=W_d^V{W}_p^V\mathbf{Y}$。其中W(·)p是1×1逐点卷积，W(·)d是3×3深度卷积。 我们在网络中使用无偏差卷积层。 接下来，我们重塑query和key投影，使它们的点积交互生成大小为 ${\mathbb{R}}^{\hat{C}\times \hat{C}}$ 的转置注意图 A，而不是大小为 ${\mathbb{R}}^{\hat{H}\hat{W}\times \hat{H}\hat{W}}$的巨大常规注意图 ,。 总体而言，MDTA 流程定义为：
$\hat{\mathbf{X}}=W_p\mathrm{Attemption}\left(\hat{\mathbf{Q}},\hat{\mathbf{K}},\hat{\mathbf{V}}\right)+\mathbf{X}$
$\text{Attention}\left(\hat{\mathbf{Q}},\hat{\mathbf{K}},\hat{\mathbf{V}}\right)=\hat{\mathbf{V}}\cdot \mathrm{Soft}\max \left(\hat{\mathbf{K}}\cdot \hat{\mathbf{Q}}/\alpha \right)$ (1)
$X$和$\hat{\mathbf{X}}$分别是输入输出特征图；$\hat{\text{Q}}\in {\mathbb{R}}^{\hat{H}\hat{W}\times \hat{C}}$ ;$\hat{\text{K}}\in {\mathbb{R}}^{\hat{C}\times \hat{H}\hat{W}}$和$\hat{\text{V}}\in {\mathbb{R}}^{\hat{H}\hat{W}\times \hat{C}}$矩阵matrices由原来的尺寸${\mathbb{R}}^{\hat{H}\times \hat{W}\times \hat{C}}$重新成形得到。 这里，α是一个可学习的缩放参数，用于在应用Softmax函数之前控制k和q的点积的大小。与传统的多头 SA 类似，我们将通道数量划分为“头”并并行学习单独的注意力图。

Gated-Dconv Feed-Forward Network(GDFN)

为了transform features，常规前馈网络 (FN) vison-transformer、Transformer分别对每个像素位置进行相同的操作。 它使用两个 1×1 卷积，一个用于扩展特征通道（通常按系数 γ=4），第二个用于将通道缩减回原始输入维度。 在隐藏层中应用非线性。 在这项工作中，我们提出了 FN 的两个基本修改以改进表示学习：（1）门控机制，和（2）深度卷积。 我们的 GDFN 架构如图 2（b）所示。 门控机制被表述为线性变换层的两个平行路径的逐元素乘积，其中一个由 GELU 非线性激活。 与 MDTA 一样，我们在GDFN中也包括深度卷积来编码来自空间相邻像素位置的信息，这有助于学习局部图像结构以进行有效的恢复。给定一个输入张量$\mathbf{X}\in {\mathbb{R}}^{\hat{H}\times \hat{W}\times \hat{C}}$，GDFN公式化为：
$\hat{\mathbf{X}}=W_p^0\mathrm{Gating}\left(\mathbf{X}\right)+\mathbf{X}$
$\mathrm{Gating}(\mathbf{X})=\varphi (W_d^1{W}_p^1(\mathrm{LN}(\mathbf{X})))\odot W_d^2{W}_p^2(\mathrm{Lal}(\mathbf{X}))$ (2)
其中，$\odot$表示元素相乘，φ表示Gelu非线性，Ln表示层归一化。 总体而言，GDFN通过我们的流水线中的各个层次控制信息流，从而允许每个层次关注与其他层次互补的细节。 也就是说，与MDTA相比，GDFN提供了一个不同的角色（专注于用上下文信息丰富特性）。 由于所提出的GDFN比常规FN运算量大，因此减小了展开比γ，使参数和计算量相近。

Progressive Learning

基于 CNN 的恢复模型通常在固定大小的图像块上进行训练。 然而，在小裁剪块上训练 Transformer 模型可能不会编码全局图像统计信息，从而在测试时提供全分辨率图像的次优性能。 为此，我们执行渐进式学习，其中网络在早期阶段接受较小图像块的训练，并在后期训练阶段接受逐渐变大的图像块训练。 通过渐进式学习在混合大小的补丁上训练的模型在测试时显示出增强的性能，其中图像可以具有不同的分辨率（图像恢复中的常见情况）。 渐进式学习策略的行为方式与课程学习过程类似，在课程学习过程中，网络从一项更简单的任务开始，逐渐转向学习更复杂的任务（需要保留精细的图像结构/纹理）。 由于在大补丁上进行训练是以延长时间为代价的，因此随着补丁大小的增加，我们减少了批量大小，以保持与固定补丁训练相似的每个优化步骤的时间。

Experiments and Analysis

我们在四个图像处理任务的基准数据集和实验设置上对所提出的Reatormer进行了评估：(a)图像去雨，(b)单图像运动去模糊，©散焦去模糊（在单图像和双像素数据上）和(d)图像去噪（在合成和真实数据上）。 关于数据集、训练参数和其他可视化结果的更多细节在补充材料中提出。 在表中，被评估方法的最佳和次佳质量分数被突出显示和下划线。 实施细节。 我们为不同的图像恢复任务训练单独的模型。 在所有实验中，我们使用以下训练参数，除非另有提及。 我们的恢复器采用4级编码器-解码器。（以下只展示去噪结果）

Image Denoising Results

我们对加性高斯白噪声生成的合成基准数据集(SET12、BSD68、Urban100
、Kodak24[20]和McMaster[104])以及真实世界的数据集（SIDD和
DND）进行去噪实验。跟biasFreeDenosing、
DRUNet一样，我们采用无偏Restormer进行去噪。

高斯去噪（灰色图像&彩色图像）

表 4 和表 5 分别显示了灰度和彩色图像去噪的几个基准数据集上不同方法的 PSNR 分数。 与现有方法 SwinIR、DPIR一致，我们在测试中包括噪声级别 15、25 和 50。 评估的方法分为两个实验类别：(1) 学习单个模型来处理各种噪声水平，以及 (2) 为每个噪声水平学习一个单独的模型。 我们的 Restormer 在不同数据集和噪声水平的实验设置下都实现了最先进的性能。 具体来说，对于高分辨率 Urban100数据集上具有挑战性的噪声级别 50，Restormer 比之前最好的基于 CNN 的方法DRUNet获得了 0.37 dB 的增益，并且比最近的基于transformer网络的SwinIR方法获得了 0.31 dB 的提升。
如表 4 所示。表 5 中的高斯颜色去噪可以观察到类似的性能提升。值得一提的是，DRUNet [需要噪声级别图作为附加输入，而我们的方法只采用噪声图像。 此外，与 SwinIR 相比，我们的 Restormer 的 FLOPs 减少了 3.14 倍，运行速度提高了 13 倍。 图 6 显示了灰度去噪（顶行）和颜色去噪（中间行）的不同方法的去噪结果。 我们的 Restormer 可恢复干净清晰的图像。

真实图像去噪

表 6 显示，我们的方法是唯一一种在两个数据集上都超过 40 dB PSNR 的方法。 值得注意的是，在 SIDD 数据集上，我们的 Restormer 分别比之前最好的 CNN 方法MIRNet和 Transformer 模型 Uformer获得了 0.3 dB 和 0.25 dB 的 PSNR 增益。 图 6（底行）显示我们的 Restormer 在不影响精细纹理的情况下生成干净的图像。

Ablation Studies(消融实验）

对于消融实验，我们在大小为 128×128 的图像块上训练高斯颜色去噪模型，仅进行 100K 次迭代。 在 Urban100 [29] 上进行了测试，并针对具有挑战性的噪声水平 σ=50 进行分析。 FLOPs 和推理时间是在 256×256 的图像大小上计算的。 表 7-10 表明我们的贡献产生了质量性能改进。 接下来，我们分别描述每个组件的影响。

Improvements in multi-head attention

表 7c 表明我们的 MDTA 提供了比基线高 0.32 dB 的有利增益（表 7a）。 此外，通过深度卷积将局部性引入 MDTA 可提高稳健性，因为移除局部性会导致 PSNR 下降（参见表 7b）。

Improvements in feed-forward network (FN)

表 7d 显示 FN 中用于控制信息流的门控机制比传统 FN 产生 0.12 dB 增益 [77]。 与多头注意力一样，将局部机制引入 FN 也带来了性能优势（见表 7e）。 我们通过结合门控深度卷积进一步加强 FN。 我们的 GDFN（表 7f）在噪声级别为 50 时实现了比标准 FN [77] 0.26 dB 的 PSNR 增益。总的来说，我们的 Transformer 块贡献导致比基线高 0.51 dB 的显着增益。

Design choices for decoder at level-1.（1级解码器的设计）

为了在第 1 层使用解码器聚合编码器功能，我们在连接操作后不使用 1×1 卷积（将通道减少一半）。 它有助于保留来自编码器的精细纹理细节，如表 8 所示。这些结果进一步证明了在细化阶段添加 Transformer 块的有效性。

Impact of progressive learning.

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

Deeper or wider Restormer?

表 10 显示，在类似的参数/FLOPs 预算下，深-窄模型比其宽-浅对应模型执行得更准确。 但是，由于并行化，更宽的模型运行得更快。 在本文中，我们使用 deep-narrow Restormer。

conclusion

我们提出了一种图像恢复 Transformer 模型 Restormer，它在处理高分辨率图像方面计算效率高。 我们将关键设计引入 Transformer 块的核心组件，以改进特征融合和提取。 具体来说，我们的 multiDconv head transposed attention (MDTA) 模块通过跨通道而不是空间维度应用自注意力来隐式建模全局上下文，因此具有线性复杂度而不是二次复杂度。 此外，所提出的门控 Dconv 前馈网络 (GDFN) 引入了一种门控机制来执行受控的特征转换。 为了将 CNN 的优势整合到 Transformer 模型中，MDTA 和 GDFN 模块都包含用于编码空间局部上下文的深度卷积。 在 16 个基准数据集上进行的大量实验表明，Restormer 在众多图像恢复任务中实现了最先进的性能。