第33篇SwinIR

Posted 2022-04-28 AI浩

SWinIR实战

论文地址：https://arxiv.org/pdf/2108.10257.pdf

预训练模型下载：https://github.com/JingyunLiang/SwinIR/releases

训练代码下载：https://github.com/cszn/KAIR

测试：https://github.com/JingyunLiang/SwinIR

摘要

图像恢复是一个长期存在的低级视觉问题，旨在从低质量图像（例如，缩小、噪声和压缩图像）中恢复高质量图像。虽然最先进的图像恢复方法是基于卷积神经网络的，但很少有人尝试使用 Transformer 进行高级视觉任务的表现令人印象深刻。在本文中，我们提出了一种基于 Swin Transformer 的强大基线模型 SwinIR 用于图像恢复。 SwinIR由浅层特征提取、深层特征提取和高质量图像重建三部分组成。特别是，深度特征提取模块由几个残差 Swin Transformer 块 (RSTB) 组成，每个残差块都有几个 Swin Transformer 层和一个残差连接。我们对三个具有代表性的任务进行了实验：图像超分辨率（包括经典、轻量级和真实世界的图像超分辨率）、图像去噪（包括灰度和彩色图像去噪）和 JPEG 压缩伪影减少。实验结果表明，SwinIR 在不同任务上的性能优于最先进的方法高达 0.14∼0.45dB，而参数总数最多可减少 67%。

一、简介

图像恢复，例如图像超分辨率（SR）、图像去噪和 JPEG 压缩伪影减少，旨在从其低质量退化的对应物重建高质量的干净图像。 由于几项革命性的工作 [18, 40, 90, 91]，卷积神经网络 (CNN) 已成为图像恢复的主要主力。

大多数基于 CNN 的方法都专注于精细的架构设计，例如残差学习 [43, 51] 和密集连接 [97, 81]。 尽管与传统的基于模型的方法[73、14、28]相比，性能得到了显着提高，但它们通常存在两个源自基本卷积层的基本问题。 首先，图像和卷积核之间的交互是与内容无关的。 使用相同的卷积核来恢复不同的图像区域可能不是最好的选择。 二、在局部处理的原理下，卷积对长距离依赖建模无效。

作为 CNN 的替代方案，Transformer [76] 设计了一种自注意力机制来捕获上下文之间的全局交互，并在多个视觉问题 [6、74、19、56] 中显示出良好的性能。 然而，用于图像恢复的视觉变换器 [9, 5] 通常将输入图像划分为固定大小的块（例如，48×48）并独立处理每个块。 这种策略不可避免地会产生两个缺点。 首先，边界像素不能利用补丁之外的相邻像素进行图像恢复。 其次，恢复的图像可能会在每个补丁周围引入边界伪影。 虽然这个问题可以通过补丁重叠来缓解，但它会带来额外的计算负担。

最近，Swin Transformer [56] 展示了巨大的前景，因为它融合了 CNN 和 Transformer 的优点。 一方面，由于局部注意力机制，它具有CNN处理大尺寸图像的优势。 另一方面，它具有 Transformer 的优势，可以使用移位窗口方案对远程依赖进行建模。

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

与流行的基于 CNN 的图像恢复模型相比，基于 Transformer 的 SwinIR 有几个优点：（1）图像内容和注意力权重之间基于内容的交互，可以解释为空间变化的卷积 [13,21,75]。 (2) 移位窗口机制启用远程依赖建模。 (3) 用更少的参数获得更好的性能。例如，如图 1 所示，与现有的图像 SR 方法相比，SwinIR 以更少的参数实现了更好的 PSNR。

二、相关工作

2.1、 图像恢复

与通常基于模型的传统图像恢复方法 [28,72,73,62,32] 相比，基于学习的方法，尤其是基于 CNN 的方法，由于其令人印象深刻的性能而变得更加流行。 他们经常从大规模配对数据集中学习低质量和高质量图像之间的映射。 自从开创性工作 SRCNN [18]（用于图像 SR）、DnCNN [90]（用于图像去噪）和 ARCNN [17]（用于 JPEG 压缩伪影减少）以来，已经提出了一系列基于 CNN 的模型来提高模型表示能力 通过使用更精细的神经网络架构设计，例如残差块 [40, 7, 88]，密集块 [81, 97, 98] 和其他 [10, 42, 93, 78, 77, 79, 50, 48,49 , 92, 70, 36, 83, 30, 11, 16, 96, 64, 38, 26, 41, 25]。 他们中的一些人利用了 CNN 框架内的注意机制，例如通道注意 [95、15、63]、非局部注意 [52、61] 和自适应补丁聚合 [100]。

2.2、视觉转换器

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

三、方法

3.1、 网络架构

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

浅层和深层特征提取。 给定一个低质量（LQ）输入$I_{LQ}\in {\mathbb{R}}^{H\times W\times C_{in}}$（H、W 和 Cin 分别是图像的高度、宽度和输入通道数），我们使用 3×3 卷积层 HSF(·) 来提取浅层特征 $F_0\in {\mathbb{R}}^{H\times W\times C}$ 为
$$F_0=H_{SF}\left(I_{LQ}\right)\text{\hspace{1em}(1)}$$

四、实验

4.1、 实验设置

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

4.2、消融研究与讨论

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

通道号、RSTB 号和 STL 号的影响。 我们展示了 RSTB 中通道数、RSTB 数和 STL 数对模型性能的影响。 分别如图 3(a)、3(b) 和 3© 所示。 可以观察到，PSNR 与这三个超参数正相关。 对于通道数，虽然性能不断提高，但参数总数呈二次增长。 为了平衡性能和模型大小，我们在其余实验中选择 180 作为通道数。 至于RSTB数和层数，性能增益逐渐趋于饱和。 我们为它们都选择了 6 以获得相对较小的模型。

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

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

4.3、图像 SR 的结果

经典形象 SR。表 2 显示了 SwinIR（中等大小）与最先进方法之间的定量比较：DBPN [31]、RCAN [95]、RRDB [81]、SAN [15]、IGNN [100]、HAN [63 ]、NLSA [61] 和 IPT [9]。可以看出，当在 DIV2K 上进行训练时，SwinIR 在几乎所有五个基准数据集上针对所有比例因子实现了最佳性能。对于比例因子 4，在 Manga109 上的最大 PSNR 增益达到 0.26dB。请注意，RCAN 和 HAN 引入了通道和空间注意，IGNN 提出了自适应补丁特征聚合，而 NLSA 基于非局部注意机制。然而，所有这些基于 CNN 的注意力机制的表现都比提出的基于 Transformer 的 SwinIR 差，这表明了所提出模型的有效性。当我们在更大的数据集（DIV2K+Flickr2K）上训练 SwinIR 时，性能进一步大幅提升（高达 0.47dB），即使 IPT 利用 ImageNet（超过 1.3 M 图像）在训练中，并且具有大量参数（115.5M）。相比之下，即使与最先进的基于 CNN 的模型（15.4∼44.3M）相比，SwinIR 的参数数量也很少（11.8M）。

至于运行时间，基于 CNN 的代表性模型 RCAN、IPT 和 SwinIR 分别需要大约 0.2、4.5 秒和 1.1 秒来测试 1024 × 1024 图像。视觉对比如图 4 所示。SwinIR 可以恢复高频细节并减轻模糊伪影，从而产生清晰自然的边缘。相比之下，大多数基于 CNN 的方法会产生模糊的图像甚至是不正确的纹理。与基于 CNN 的方法相比，IPT 生成更好的图像，但它存在图像失真和边界伪影。

轻量级图像 SR。 我们还提供了 SwinIR（小尺寸）与最先进的轻量级图像 SR 方法的比较：CARN [2]、FALSR-A [12]、IMDN [35]、LAPAR-A [44] 和 LatticeNet [57 ]。 除了 PSNR 和 SSIM，我们还报告了参数总数和乘法累加操作（在 1280×720 的 HQ 图像上评估），以比较不同模型的模型大小和计算复杂度。 如表 3 所示，SwinIR 在不同的基准数据集上以高达 0.53dB 的 PSNR 幅度优于竞争方法，具有相似的参数总数和乘法累加操作。 这表明 SwinIR 架构对于图像恢复非常有效。

**真实世界的图像 SR。**图像 SR 的最终目标是用于实际应用。最近，张等人[89]为真实世界的图像 SR 提出了一个实用的退化模型 BSRGAN，并在真实场景中取得了令人惊讶的结果。为了测试 SwinIR 在现实世界 SR 中的性能，我们使用与 BSRGAN 相同的退化模型重新训练 SwinIR，以进行低质量图像合成。由于没有真实的高质量图像，我们仅提供与代表性双三次模型 ESRGAN [81] 和最先进的真实世界图像 SR 模型 RealSR [37]、BSRGAN [89] 和 Real-ESRGAN 的视觉比较[80]。如图 5 所示，SwinIR 产生具有清晰锐利边缘的视觉上令人愉悦的图像，而其他比较方法可能会出现不令人满意的伪影。此外，为了充分利用 SwinIR 在实际应用中的潜力，我们进一步提出了一个大型模型并在更大的数据集上对其进行训练。实验表明，与当前模型相比，它可以处理更复杂的损坏，并在真实世界的图像上取得更好的性能。由于页数限制，详情请见我们的项目页面https://github.com/JingyunLiang/SwinIR。

4.4. JPEG 压缩伪影减少的结果

表 4 显示了 SwinIR 与最先进的 JPEG 压缩伪影减少方法的比较：ARCNN [17]、DnCNN-3 [90]、QGAC [20]、RNAN [96]、RDN [98] 和 DRUNet [88] . 所有比较的方法都是基于 CNN 的模型。 在 [98, 88] 之后，我们在两个基准数据集（Classic5 [22] 和 LIVE1 [67]）上测试了 JPEG 质量因子 10、20、30 和 40 的不同方法。正如我们所见，所提出的 SwinIR 具有平均 PSNR 增益 在不同品质因数的两个测试数据集上至少为 0.11dB 和 0.07dB。 此外，与之前最好的模型 DRUNet 相比，SwinIR 只有 1150 万个参数，而 DRUNet 是一个大模型，有 3270 万个参数。

4.5.图像去噪结果

我们分别在表 5 和表 6 中显示灰度和彩色图像去噪结果。比较的方法包括传统模型 BM3D [14] 和 WNNM [29]、基于 CNN 的模型 DnCNN [90]、IRCNN [91]、FFDNet [92]、N3Net [65]、NLRN [52]、FOCNet [38]、RNAN [96]、MWCNN [54] 和 DRUNet [88]。在 [90, 88] 之后，比较的噪声水平包括 15、25 和 50。可以看出，我们的模型比所有比较的方法实现了更好的性能。特别是，它在具有 100 个高分辨率的大型 Urban100 数据集上超过了最先进的模型 DRUNet 高达 0.3dB 测试图像。值得指出的是，SwinIR 只有 12.0M 参数，而 DRUNet 有 32.7M 参数。这表明 SwinIR 架构在学习特征表示以进行恢复方面非常有效。不同方法的灰度和彩色图像去噪的视觉比较如图 6 和图 7 所示。我们可以看到，我们的方法可以去除严重的噪声破坏并保留高频图像细节，从而产生更清晰的边缘和更自然的纹理。相比之下，其他方法要么过于平滑，要么过于锐利，无法恢复丰富的纹理。

五、 结论

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