图像复原 ---SwinIR

系列文章目录

---

论文名称：Image Restoration Using Swin Transformer
论文地址：http://arxiv.org/abs/2108.10257
代码地址：https://github.com/jingyunliang/swinir
发表时间：2021
引用邻域：图像超分辨、图像去噪（本文仅关注去噪部分）
主要模块：浅层特征提取、深层特征提取、重建模块

---

---

摘要

在本文中，我们提出了一个强大的基线模型SwinIR，用于基于Swin Transformer的图像恢复。SwinIR由三部分组成：浅层特征提取、深层特征提取和高质量图像重建。特别地，深度特征提取模块由几个残差Swin Transformer块（RSTB)组成，每个有几个Swin Transformer层和一个剩余连接。我们对三个代表性任务进行了实验：图像超分辨率（包括经典、轻量级和真实世界图像超分辨像超分辨率）、图像去噪（包括灰度和彩色图像去噪）和JPEG压缩伪影减少。实验结果表明，SwinIR在不同任务上由于最先进的方法高达0.14 -0.45dB，而参数总数最多可减少67%。

Introduction

CNN balabala
卷积存在的两个问题：1.使用相同的卷积核来恢复不同的图像区域可能不是最好的选择。2.在局部处理的原则下，卷积对于远程依赖建模是无效的。

当下Transformer模型存在的问题：1.边界像素不能利用补丁外的相邻像素进行图像恢复。2.恢复后的图像可能会在每个补丁周围引入边界伪影。虽然两个问题可以通过补丁重叠来缓解，但它会引入额外的计算负担

Swin Transformer优势：1.由于局部注意机制，它具有CNN处理大尺寸图像的优势。2.它具有Transformer的优势，可以使用移位窗口方案对远程依赖进行建模。

在本文中，我们提出了一种基于Swin Transformer的图像恢复模型，即swinIR.更具体地说，swinIR由三个模块组成：浅层特征提取、深度特征提取和高质量图像重建模块。浅层特征提取模块使用卷积层提取浅层特征，直接传递给重建模块，以保留低频信息。深度特征提取模块主要由残差Swin Transformer块（RSTB）组成，每个块都利用多个Swin Transformer层进行局部注意和跨窗口交互。此外，我们在块的末尾添加一个卷积层以增强特征，并使用残差连接为特征聚合提供捷径。最后，在重建模块中融合浅层和深层特征以进行高质量图像重建。

与流行的基于CNN的图像恢复模型相比，基于transformer的SwinIR有几个好处：1.图像内容和注意力权重之间基于内容交互，可以解释为空间变化的卷积attention cnn、paddleVit。2.移动窗口机制支持远程依赖建模。3.用更少的参数获得更好的性能。

Realated work（相关工作）

Image Restoration（图像复原）

与通常基于模型的传统图像恢复方法相比，基于学习的方法，尤其是基于CNN的方法，由于其令人印象深刻的性能而变得更受欢迎。它们经常从大规模配对数据集中学习低质量和高质量图像之间的映射。自从开创性工作SRCNN（用于图像SR)、DnCNN（用于图像去噪）、和ARCNN（用于JPEG压缩伪影减少）以来，一系列基于CNN的模型已经出现建议通过使用更精细的神经网络架构设计来提高模型表示能力，例如残差块DPIR、密集块RDN和其他FDDNet、RNAN。其中一些利用了CNN框架内的注意力机制，例如通道注意力、非局部注意力NLRN和自适应补丁聚合。

Vision transformer(ViT)

最近，自然语言处理模型Transformer在计算机视觉社区中广受欢迎。当用于图像分类、目标检测、分割和人群计数，它通过探索不同区域之间的全局交互来学习关注重要的图像区域地区。由于其令人印象深刻的性能，Transformer也被引入用于图像恢复。陈等。提出了一种基于标准Transformer的用于各种恢复问题的骨干模型IPT。然而，IPT依赖于大量参数（超过115.5M参数）、大规模数据集（超过1.1M图像）和多任务学习以获得良好的性能。曹等提出了VSR-Transformer，它使用自注意力机制在视频 SR 中更好地融合特征，但图像特征仍然是从 CNN 中提取的。 此外，IPT 和 VSR-Transformer 都是 patch-wise attention，可能不适用于图像恢复。 此外，一项并行工作提出了一种基于 Swin Transformer的 U 形架构

Method

Network architecture(网络架构）

如图2所示，SwinIR由三个模块组成：浅层特征提取、深度特征提取和高质量（HQ）图像重建模块。 我们对所有恢复任务使用相同的特征提取模块，但对不同的任务使用不同的重建模块。

Shallow and deep feature extraction（浅层深层特征提取）

给定一张低质量（LQ)$I_{LQ}\in {\mathbb{R}}^{H\times W\times C_{in}}$作为输入(H、W和$C_{in}$分别是图像高度、宽度和输入通道数），我们使用3*3卷积层$H_{SF}(\cdot )$从输入低质量的图像中提取浅层特征$F_0\in {\mathbb{R}}^{H\times W\times C}$
$F_0=H_{SF}(I_{LQ})$ （1）
其中c是特征通道数。卷积层擅长早期视觉处理，使得更稳定的优化和更好的结果。它还提供了一种将输入图像空间映射到更高维特征空间的简单方法。然后我们从$F_0$中提取更深层次的特征$F_{DF}\in {\mathbb{R}}^{H\times W\times C}$。
$F_{DF}=H_{DF}(F_0)$ (2)
其中$H_{DF}(\cdot )$是深度特征提取模块，它包含k个残差Swin Transformer块（RSTB）和一个3 * 3卷积层。更具体地说，中间特征F1，F2，…，$F_k$和输出的深度特征$F_{DF}$逐块提取为
$F_i=H_{RST{B}_i}(F_{i-1}),i=1,2,\dots ,K$
$F_{DF}=H_{CONV}(F_K)$ (3)
其中，HRSTBI(·)表示第i个RSTB，HCONV是最后一个卷积层.在特征提取的最后一层用一个卷积可以将卷积操作的归纳偏差带入基于Transformer的网络中，为后期浅层和深层的聚合打下更好的基础。

Image reconstruction（图像重建）

以SR（图像超分）为例，我们通过聚合浅层和深层特征来重建高质量图像$I_{RHQ}$
$I_{RHQ}=H_{REC}(F_0+F_{DF})$ (4)
其中$H_{REC}$(·)是重构模块的函数。 浅层特征主要包含低频特征，而深层特征则侧重于恢复丢失的高频特征。 SwinIR通过较长的跳过连接，将低频信息直接传递给重构模块，帮助深度特征提取模块聚焦高频信息，稳定训练。 对于重构模块的实现，我们使用亚像素卷积层对特征进行上采样。
对于图像去噪、JPEG压缩伪影降低等不需要上采样的任务，采用单个卷积层进行重建。 此外，我们利用残差学习来重建LQ图像与HQ图像之间的残差，而不是HQ图像。 这表示为
$I_{RHQ}=H_{SwinIR}(I_{LQ})+I_{LQ}$ (6)
其中$H_{swinIR}$(·)表示swinIR的函数。

Loss funcion（损失函数）

对于图像去噪和JPEG压缩伪影减少，我们使用Charbonnier损失
$\mathcal{L}=\sqrt{\Vert I_{RHQ}-I_{HQ}{\Vert }^2+{ϵ}^2}$ (7)
其中$ϵ$是根据经验设置为$10^{-3}$的常数

Residual Swin Transformer Block (RSTB)

如图2（a）所示，残差Swin Transformer块（RSTB)是一个带有Swin Transformer层（STL）和卷积层的残差块。给定第i个RSTB的输入特征$F_{i,0}$,我们首先提取中间特征$F_{i,1}$、$F_{i,2}$…$F_{i,L}$,通过L Swin Transformer层为
$F_{i,j}=H_{ST{L}_{i,j}}(F_{i,j-1}),j=1,2,\dots ,L$ (8)
$H_{ST{L}_{i,j}}(\cdot )$卷积层是第i个RSTB的第j个Swin Transformer层。然后我们在残差连接前加一个卷积层。RSTB输出的公式如下：
$F_{i,out}=H_{CON{V}_i}(F_{i,L})+F_{i,0}$ (9)
$H_{CON{V}_i}(\cdot )$是第i个RSTB的卷积层。这种设计有两个好处。首先，虽然，Transformer可以被视为空间变化卷积的特定实例，但具有空间不变滤波器的卷积层可以增强SwinIR的平移等变性。其次，残差连接提供了从不同块到重建模块的基于身份的连接，允许聚合不同级别的特征。

Swin Transformer layer

Swin Transformer层是基于标准的多头注意力的原始Transformer层 。主要的不同之处在于局部注意力和移动窗口机制。如图2（b）所示。给定输入的大小HWC，Swin Transformer首先将输入输入重塑为通过M*M的非局部窗口将输入划分为$\text{a}\frac{HW}{M^2}\times M^2\times C$的特征图，$\frac{HW}{M^2}$是所有的窗口数量。然后，它分别为每个窗口计算标准自注意力（即局部注意力）。对于局部窗口特征$X\in {\mathbb{R}}^{M^2\times C}$，query、key、value计算为：
$Q=X{P}_Q,K=X{P}_K,V=X{P}_V$ (10)
其中，$P_Q$、$P_K$、$P_V$是在不同窗口之间共享的投影矩阵。一般来说，我们有$Q,K,V\in {\mathbb{R}}^{M^2\times d}$。因此，注意力矩阵有局部窗口种的自注意力机制计算为：
$\text{Attenion}(Q,K,V)=\text{SoftMax}(Q{K}^T/\sqrt{d}+B)V$ (11)
其中 B 是可学习的相对位置编码。 在实践中，按照[76]，我们并行执行 h 次注意力函数，并将结果连接起来用于多头自注意力 (MSA)。 接下来，多层感知器 (MLP) 具有两个完全连接的层，它们之间具有 GELU 非线性，用于进一步的特征转换。 在 MSA 和 MLP 之前添加了 LayerNorm (LN) 层，并且两个模块都使用了残差连接。 整个过程表述为
$X=\mathrm{MSA}(\mathrm{LN}(X))+X$
$X=\mathrm{Mlp}(\mathrm{LN}(X))+X$ (12)
然而，当不同层的分区固定时，局部窗口之间没有连接。 因此，交替使用常规和移位窗口分区来启用跨窗口连接，其中移位窗口分区意味着在分区之前将特征移动$(\lfloor \frac{M}{2}\rfloor ,\lfloor \frac{M}{2}\rfloor )$像素。

Experiments

Experimental Setup(实验设置）

对于经典图像SR、真实世界图像SR、图像去噪和JPEG压缩伪影减少，RSTB数、STL数、窗口大小、通道数和注意力头数一般分别设置为6、6、8、180和6 . 一个例外是窗口大小设置为 7 以减少 JPEG 压缩伪影，因为我们观察到使用 8 时性能显着下降，这可能是因为 JPEG 编码使用 8 × 8 图像分区。对于轻量级图像 SR，我们将 RSTB 数量和通道数量分别减少到 4 和 60。 当在测试中使用自集成策略时，我们用符号“+”标记模型，例如 SwinIR+。 由于篇幅限制，培训和评估细节在补充中提供。

Ablation Study and Discussion

对于消融研究，我们在 DIV2K上为经典图像 SR (×2) 训练 SwinIR，并在 Manga109上对其进行测试。

Impact of channel number, RSTB number and STL number（通道数、RSTB数和STL数的影响）

我们给出了RSTB中的信道数、RSTB数和STL数对模型性能的影响。 如图3(a)、3(b)和3©所示。 观察到PSNR与这三个超参数呈正相关。 对于通道数，虽然性能不断增加，但参数总数以二倍的速度增长。 为了平衡性能和模型大小，我们在其余实验中选择180作为通道数。 对于RSTB数和层数，性能增益逐渐饱和。 我们对两者都选择6，以获得一个相对较小的模型。

Impact of patch size and training image number; modelconvergence comparison（补丁大小和训练图像数量的影响； 模型收敛性比较）

我们将所提出的 SwinIR 与代表性的基于 CNN 的模型 RCAN 进行比较，以比较基于 Transformer 和基于 CNN 的模型的差异。 从图 3（d）可以看出，SwinIR 在不同的 patch size 上表现优于 RCAN，并且当 patch size 较大时 PSNR gain 变大。 图 3(e) 显示了训练图像数量的影响。 当百分比大于 100%（800 张图像）时，使用来自 Flickr2K 的额外图像进行训练。 有两个观察结果。 首先，正如预期的那样，SwinIR 的性能随着训练图像数量的增加而提高。 其次，与 IPT 中观察到的基于 Transformer 的模型严重依赖大量训练数据不同，SwinIR 比使用相同训练数据的基于 CNN 的模型取得了更好的结果，即使数据集很小（即 25%， 200 张图像）。 我们还在图 3（f）中绘制了 SwinIR 和 RCAN 训练期间的 PSNR。 很明显，SwinIR 比 RCAN 收敛得更快更好，这与之前的观察结果相矛盾，即基于 Transformer 的模型通常会出现模型收敛速度慢的问题。

Impact of residual connection and convolution layer in RSTB（RSTB中残差连接和卷积层的影响）

表1给出了RSTB中的四种残差连接变体：无残差连接、使用1×1卷积层、使用3×3卷积层和使用三个3×3卷积层（中间层的通道数设置为网络通道的四分之一） 数字）。 从表中，我们可以有以下观察结果。 首先，RSTB 中的残余连接很重要，因为它将 PSNR 提高了 0.16dB。 其次，使用 1 × 1 卷积带来的改进很小，可能是因为它不能像 3×3 卷积那样提取局部邻近信息。 第三，虽然使用三个 3×3 卷积层可以减少参数数量，但性能略有下降。

Results on Image Denoising

我们分别在表5、表6种显示了灰度图像和彩色图像去噪的结果。同传统模型BM3D、WNNM，基于cnn的模型 DnCNN、IRCNN、FFDNet、N3Net、NLRN、FOCNet、RNAN、MWCNN和DRUNet相比较，比较从噪声水平包括15、25和50.可以看到，我们的模型比所有比较方法都获得了更好的性能。特别是，在具有100张高分辨率测试图像的大型Urban100数据集上，它超过了最先进的DRUN模型0.3dB。值得指出的是，SwinIR只有12.0M个参数，而DRUN有32.7M个参数。这表明SwinIR体系结构在学习用于恢复的特征表示方面是高效的。图6和图7显示了不同方法的灰度和彩色图像去噪的视觉比较。正如我们所看到的，我们的方法可以去除严重的噪声破坏，并保留高频图像细节，从而获得更清晰的边缘和更自然的纹理。相比之下，其他方法要么过于平滑，要么过于清晰，无法恢复丰富的纹理。

Conclusion

在本文中，我们提出了一种基于Swin Transformer的图像恢复模型SwinIR。该模型由三部分组成：浅层特征提取、深层特征提取和HR（高分辨率图像）重建模块。 特别是，我们使用一堆残差 Swin Transformer 块 (RSTB) 进行深度特征提取，每个 RSTB 由 Swin Transformer 层、卷积层和残差连接组成。 大量实验表明，SwinIR 在三个代表性图像恢复任务和六个不同设置上实现了最先进的性能：经典图像 SR、轻量级图像 SR、真实世界图像 SR、灰度图像去噪、彩色图像去噪和 JPEG 压缩伪影 减少，这证明了所提出的 SwinIR 的有效性和普遍性。 将来，我们会将模型扩展到其他恢复任务，例如图像去模糊和去雨。