《SwinIR：使用Swin-Transformer图像恢复》学习笔记

paper：2108.10257

GitHub：GitHub - JingyunLiang/SwinIR： SwinIR：使用 Swin Transformer 进行图像修复（官方仓库）

摘要

1、Introduction

2、Related Work

2.1 图像修复

2.2 视觉Transformer

3、方法

3.1 网络框架

3.1.1 浅层与深层特征提取

3.1.2 图像重建

3.1.3 损失函数

3.2 残差Swin Transformer块（RSTB）

3.2.1. Swin Transformer层（STL）

4、实验

4.1 实验设置

4.2 消融实验

4.2.1 通道数、RSTB数量和STL数量的影响

编辑 4.2.2 patch大小与训练图像数量的影响

编辑 4.2.3 RSTB中残差连接和卷积层的影响

编辑 4.3 图像超分辨率结果

Lightweight Image SR

Real-world Image SR

4.4 JPEG压缩伪影去除

4.5 图像去噪

摘要

图像修复是一个长期存在的低级视觉问题，旨在从低质量图像（例如，降采样、噪声污染和压缩的图像）中恢复高质量图像。虽然最先进的图像修复方法大多基于卷积神经网络（CNN），但很少有研究尝试使用在高级视觉任务中表现出色的Transformer模型。在本文中，我们提出了一种基于Swin Transformer的强基线模型——SwinIR，用于图像修复。SwinIR包含三个部分：浅层特征提取、深层特征提取和高质量图像重建。特别地，深层特征提取模块由若干残差Swin Transformer块（Residual Swin Transformer Block，RSTB）组成，每个块包含若干Swin Transformer层以及一个残差连接。我们在三种具有代表性的任务上进行了实验：图像超分辨率（包括经典、轻量级和真实场景图像超分辨率）、图像去噪（包括灰度图像去噪和彩色图像去噪）以及JPEG压缩伪影消除。实验结果表明，SwinIR在不同任务中均优于现有最先进方法，在性能上提升了0.14到0.45 dB，同时参数总量最多可减少67%。

1、Introduction

图像修复（如图像超分辨率（SR）、图像去噪和JPEG压缩伪影消除）旨在从低质量的退化图像中重建高质量的干净图像。自从一些革命性工作推出后，卷积神经网络（CNN）已成为图像修复的主要工具。大多数基于CNN的方法专注于精细的架构设计，例如残差学习和密集连接。尽管与传统基于模型的方法相比性能显著提高，但它们通常面临由基础卷积层引起的两个基本问题。

首先，图像与卷积核之间的交互是内容无关的。使用相同的卷积核恢复不同的图像区域可能并非最佳选择。其次，在局部处理的原则下，卷积难以有效建模长距离的依赖关系。

作为CNN的替代方案，Transformer设计了一个自注意力机制以捕获上下文之间的全局交互，并在若干视觉问题中表现出令人期待的性能。然而，用于图像修复的视觉Transformer通常将输入图像划分为固定大小（如48×48）的块，并独立处理每个块。这种策略不可避免地导致两个缺点。首先，边缘像素无法利用块外的邻域像素进行图像修复。其次，修复后的图像可能会在每个块的边缘引入伪影。虽然可以通过块重叠来缓解这一问题，但这会增加额外的计算负担。

最近，Swin Transformer显示出了巨大的潜力，因为它结合了CNN和Transformer的优点。一方面，它利用局部注意力机制，继承了CNN在处理大尺寸图像上的优势；另一方面，通过平移窗口（shifted window）机制，它利用了Transformer在建模长距离依赖关系上的能力。

在本文中，我们提出了一种基于Swin Transformer的图像修复模型，即SwinIR。具体来说，SwinIR由三个模块组成：浅层特征提取模块、深层特征提取模块和高质量图像重建模块。浅层特征提取模块通过卷积层提取浅层特征，并将其直接传递到重建模块，以保留低频信息。深层特征提取模块主要由残差Swin Transformer块（RSTB）组成，每个块利用多个Swin Transformer层进行局部注意力建模和跨窗口交互。此外，我们在块的末端添加了一个卷积层用于特征增强，并通过残差连接为特征聚合提供了一条快捷路径。最终，浅层特征和深层特征在重建模块中融合，以实现高质量的图像重建。

与流行的基于CNN的图像修复模型相比，基于Transformer的SwinIR具有以下几个优势：
(1) 图像内容与注意力权重之间基于内容的交互，可被解释为空间可变卷积；
(2) 平移窗口机制使得长距离依赖关系的建模成为可能；
(3) 在参数量更少的情况下实现了更好的性能。例如，正如图所示，SwinIR在参数量更少的情况下，较现有的图像超分辨率方法取得了更高的PSNR。

2、Related Work

2.1 图像修复

与传统基于模型的图像修复方法相比，基于学习的方法，特别是基于CNN的方法，由于其令人瞩目的性能，已变得更加流行。这些方法通常通过大规模配对数据集学习低质量图像与高质量图像之间的映射关系。自从开创性工作SRCNN（用于图像超分辨率）、DnCNN（用于图像去噪）和ARCNN（用于JPEG压缩伪影消除）提出以来，许多基于CNN的模型相继被提出，这些模型通过更精细的神经网络架构设计（如残差块、密集块等）来提升模型的表征能力。

此外，一些方法在CNN框架中引入了注意力机制，例如通道注意力、非局部注意力和自适应块聚合等。这些改进进一步提高了模型在图像修复任务中的性能。

2.2 视觉Transformer

近年来，自然语言处理模型Transformer在计算机视觉领域受到广泛关注。在图像分类、目标检测、分割和人群计数等任务中，Transformer通过探索不同区域之间的全局交互关系，学会关注图像中的重要区域。由于其卓越的性能，Transformer也被引入到图像修复任务中。

Chen等人提出了一种基于标准Transformer的主干模型IPT，用于多种图像修复任务。然而，IPT依赖于大量的参数（超过115.5M参数）、大规模数据集（超过110万张图像）以及多任务学习来实现其性能。Cao等人提出了VSR-Transformer，该方法在视频超分辨率任务中利用自注意力机制实现了更好的特征融合，但其图像特征仍然是从CNN中提取的。此外，无论是IPT还是VSR-Transformer，其注意力机制都是基于图像块的，这可能并不完全适用于图像修复任务。

此外，另一项同期工作提出了一种基于Swin Transformer的U型架构，用于图像修复任务。

3、方法

3.1 网络框架

如图所示，SwinIR包含三个模块：浅层特征提取、深层特征提取以及高质量（HQ）图像重建模块。对于所有修复任务，我们使用相同的特征提取模块，但根据不同任务选择不同的重建模块。

3.1.1 浅层与深层特征提取

给定一个低质量（LQ）输入 $I_{LQ} \in \mathbb{R}^{H \times W \times C_{in}}$ （其中 $H$ 、 $W$ 和 $C_{in}$ 分别表示图像的高度、宽度和输入通道数），我们使用一个 $3 \times 3$ 的卷积层 $H_{SF}(\cdot)$ 提取浅层特征 $F_0 \in \mathbb{R}^{H \times W \times C}$ ：

$F_0 = H_{SF}(I_{LQ})$

其中 $C$ 是特征通道数。卷积层在早期视觉处理中表现出色，能够带来更稳定的优化效果和更优的结果。同时，它还能将输入图像空间映射到更高维度的特征空间。随后，我们从 $F_0$ 中提取深层特征 $F_{DF} \in \mathbb{R}^{H \times W \times C}$ ：

$F_{DF} = H_{DF}(F_0)$

如图所示， $H_{DF}(\cdot)$ 是深层特征提取模块，包含 $K$ 个残差Swin Transformer块（RSTB）和一个 $3 \times 3$ 的卷积层。

具体来说，逐块提取的中间特征 $F_1, F_2, \dots, F_K$ 以及输出的深层特征 $F_{DF}$ 的计算公式为：

$F_i = H_{RSTB_i}(F_{i-1}), \quad i=1,2,\dots,K$

$F_{DF} = H_{CONV}(F_K)$

其中 $H_{RSTB_i}(\cdot)$ 表示第 $i$ 个 RSTB（多个模块堆叠）， $H_{CONV}(\cdot)$ 是最后的卷积层。在特征提取的末尾使用卷积层可以将卷积操作的归纳偏置引入基于Transformer的网络，并为浅层与深层特征的后续聚合奠定更好的基础。

3.1.2 图像重建

以图像超分辨率（SR）为例，我们通过聚合浅层和深层特征重建高质量图像 $I_{RHQ}$ ：

$I_{RHQ} = H_{REC}(F_0 + F_{DF})$

其中 $H_{REC}(\cdot)$ 是重建模块的函数。浅层特征主要包含低频信息，深层特征则关注恢复丢失的高频信息。通过长跳跃连接，SwinIR能够将低频信息直接传递给重建模块，这有助于深层特征提取模块专注于高频信息，并稳定训练过程。在重建模块的实现中，我们使用子像素卷积层对特征进行上采样。而对于不需要上采样的任务（如图像去噪和JPEG压缩伪影去除），我们使用单个卷积层完成重建。此外，我们采用残差学习的方式来重建LQ图像与HQ图像之间的残差，而不是直接预测HQ图像：

$I_{RHQ} = H_{SwinIR}(I_{LQ}) + I_{LQ}$

其中 $H_{SwinIR}(\cdot)$ 表示SwinIR的函数。

子像素卷积层对特征进行上采样是一种用于提高图像分辨率的技术，主要用于图像超分辨率任务。其核心思想是通过卷积操作生成一个通道数为放大倍数平方的特征图，然后利用一种称为PixelShuffle的操作，将这些通道重新排列成更高分辨率的图像。

具体来说，子像素卷积层的工作流程如下：

卷积操作：输入一个低分辨率的特征图，通过卷积操作生成一个通道数为upscale_factor²的特征图，其中upscale_factor是图像放大的倍数。

PixelShuffle操作：将生成的特征图进行周期性重排，将每个像素的位置扩展为upscale_factor × upscale_factor的块，从而实现图像的放大。

3.1.3 损失函数

对于图像超分辨率（SR），通过最小化 $L_1$ 像素损失优化SwinIR的参数：

$\mathcal{L} = \|I_{RHQ} - I_{HQ}\|_1$

其中 $I_{RHQ}$ 是以 $I_{LQ}$ 为输入通过SwinIR获得的输出，也就是重建的高质量图像， $I_{HQ}$ 是对应的高质量真实图像。对于经典和轻量化的图像SR任务，我们仅使用简单的 $L_1$ 像素损失以证明网络的有效性。而对于真实场景的图像SR，我们结合像素损失、GAN损失和感知损失以提高视觉质量。

对于图像去噪和JPEG压缩伪影去除，我们使用Charbonnier损失：

$\mathcal{L} = \sqrt{\|I_{RHQ} - I_{HQ}\|_2^2 + \epsilon^2}$

其中 $\epsilon$ 是常数，经验值设为 $10^{-3}$ 。

3.2 残差Swin Transformer块（RSTB）

如图所示，RSTB 是一种结合了Swin Transformer层和卷积层的残差结构。

给定第 $i$ 个RSTB的输入特征 $F_{i}^{0}$ ，通过 $L$ 个Swin Transformer层提取中间特征 $F_{i}^{1}, F_{i}^{2}, \dots, F_{i}^{L}$ 。具体公式为：

$F_{i}^{j} = H_{STL}^{ij}(F_{i}^{j-1}), \quad j = 1, 2, \dots, L$

其中 $H_{STL}^{ij}$ 表示第 $i$ 个RSTB中的第 $j$ 个Swin Transformer层。

随后，在特征 $F_{i}^{L}$ 后加入一个卷积层，并通过残差连接得到输出特征 $F_{i}^{\text{out}}$

$F_{i}^{\text{out}} = H_{\text{CONV}}^{i}(F_{i}^{L}) + F_{i}^{0}$

其中 $H_{\text{CONV}}^{i}$ 是第 $i$ 个RSTB中的卷积操作。

这种设计的两大优势：

卷积层通过空间不变的滤波器增强了SwinIR的平移等变性；
残差连接为从不同块到重建模块提供了基于身份的连接，便于聚合不同层次的特征。

3.2.1. Swin Transformer层（STL）

Swin Transformer层基于原始Transformer的多头自注意力（Multi-head Self-Attention, MSA），但采用了局部注意力和移动窗口机制进行改进。

给定一个输入特征 $X$ 的尺寸为 $H \times W \times C$ ，Swin Transformer首先将其分割为多个 $M \times M$ 的非重叠局部窗口，每个窗口的特征维度为 $M^2 \times C$ 。然后，对每个窗口单独计算标准的自注意力（local attention）。
对于局部窗口特征 $X \in \mathbb{R}^{M^2 \times C}$ ，查询（Q）、键（K）和值（V）矩阵的计算公式为：

$Q = X P_Q, \quad K = X P_K, \quad V = X P_V$

其中 $P_Q$ 、 $P_K$ 和 $P_V$ 是投影矩阵，通常 $Q, K, V \in \mathbb{R}^{M^2 \times d}$ 。
随后，通过局部窗口内的自注意力机制计算注意力矩阵：

$\text{Attention}(Q, K, V) = \text{SoftMax}\left(\frac{Q K^T}{\sqrt{d}} + B\right)V$

其中 $B$ 是可学习的相对位置编码。

为进一步捕获更多的特征信息，采用了多头自注意力（MSA），即对注意力机制进行 $h$ 次并行计算并将结果拼接；引入多层感知机（MLP），包括两个全连接层，中间使用GELU激活函数。

每个Swin Transformer层包括两个残差连接和两个LayerNorm（LN）操作，整体流程为：

$X' = \text{MSA}(\text{LN}(X)) + X$

$X'' = \text{MLP}(\text{LN}(X')) + X'$

在固定分区下，相邻窗口之间无法建立联系。为此，Swin Transformer交替使用常规窗口分区和移动窗口分区，移动窗口分区通过在分区前将特征平移 $(\frac{M}{2}, \frac{M}{2})$ 像素，从而实现跨窗口的连接。

4、实验

4.1 实验设置

在处理经典图像超分辨率（SR）、真实世界图像超分辨率、图像去噪以及JPEG压缩伪影去除任务时，实验中采用的超参数配置为：

RSTB（Residual Swin Transformer Block）数量设为6；
STL（Swin Transformer Layer）数量设为6；
窗口大小为8；
通道数为180；
注意力头数为6。

在JPEG压缩伪影去除任务中，由于JPEG编码采用8×8图像分块，使用窗口大小为8会导致性能显著下降，因此窗口大小调整为7。

对于轻量级的图像超分辨率任务，RSTB的数量和通道数分别减小为4和60。测试时，若采用自集成策略（self-ensemble strategy），标记模型为“SwinIR+”。

4.2 消融实验

4.2.1 通道数、RSTB数量和STL数量的影响

实验表明，PSNR 随着这三个超参数的增加而提升。

PSNR（Peak Signal-to-Noise Ratio，峰值信噪比）是一种用于衡量图像或视频质量的常用客观指标，通常用于评估压缩、去噪或重建算法的性能。PSNR 值越高，表示处理后的图像与原始图像越接近，质量越高。

具体而言：

通道数的增加会使模型性能提升，但参数数量呈平方增长，因此在性能与模型大小之间权衡，通道数选择为180。
RSTB数量和STL层数的增加带来的性能增益逐渐趋于饱和，因此选择6作为最终配置。

4.2.2 patch大小与训练图像数量的影响

通过与CNN模型RCAN的对比，结果显示：

SwinIR在不同补丁大小上均优于RCAN，且补丁越大，性能提升越显著。
增加训练图像数量时，SwinIR的性能持续提升，且即使在小规模数据集（200张图像）下，SwinIR也能优于RCAN，展现了Transformer模型在有限数据条件下的优势。
从训练曲线看，SwinIR相比RCAN具有更快的收敛速度与更优的收敛效果，突破了Transformer模型通常收敛较慢的观察结论。

4.2.3 RSTB中残差连接和卷积层的影响

针对残差连接和卷积层的不同设计，实验结果表明：

残差连接对PSNR提升重要，能够带来0.16 dB的性能增益。
使用1×1卷积的性能提升有限，因为其无法像3×3卷积那样提取局部邻域信息。
使用三个3×3卷积虽然可以减少参数数量，但性能略有下降。

4.3 图像超分辨率结果

在经典图像超分辨率（SR）任务中，表中展示了 SwinIR（中型版本）与当前最先进方法的定量比较，包括 DBPN、RCAN、RRDB、SAN、IGNN、HAN、NLSA 和 IPT。实验结果显示，当使用 DIV2K 数据集进行训练时，SwinIR 在几乎所有五个基准数据集上都取得了最好的表现，并且在 Manga109 数据集上，对于放大比例为 4 时，PSNR 增益最大可达到 0.26dB。尽管 RCAN 和 HAN 引入了通道和空间注意力机制，IGNN 提出了自适应补丁特征聚合方法，NLSA 基于非局部注意力机制，但所有这些基于 CNN 的注意力机制都不如提出的基于 Transformer 的 SwinIR 模型。这表明，SwinIR 模型在图像超分辨率任务中具有较好的效果。

当我们在一个更大的数据集（DIV2K + Flickr2K）上训练 SwinIR 时，性能有了大幅提升（最大可达 0.47dB），即使 IPT 使用了超过 1.3M 图像的 ImageNet 数据集进行训练并且具有巨大的参数量（115.5M），SwinIR 也超越了 IPT 在准确度上的表现。与此相比，SwinIR 的参数量较小（11.8M），即使与最先进的基于 CNN 的模型（15.4M - 44.3M）相比，参数量也相对较少。关于运行时间，代表性的 CNN 基础模型 RCAN、IPT 和 SwinIR 在对 1024×1024 图像进行测试时，分别需要约 0.2 秒、4.5 秒和 1.1 秒。

视觉比较显示，SwinIR 可以恢复高频细节并缓解模糊伪影，从而获得清晰自然的边缘。而大多数基于 CNN 的方法则生成了模糊的图像，甚至在纹理上出现了错误。尽管 IPT 在图像质量上比基于 CNN 的方法更好，但它仍然存在图像畸变和边缘伪影。

Lightweight Image SR

我们还提供了 SwinIR（小型版本）与当前最先进的轻量级图像超分辨率方法的比较，包括 CARN、FALSR-A、IMDN、LAPAR-A 和 LatticeNet。除 PSNR 和 SSIM 外，我们还报告了不同模型的参数总数和乘加操作数（在 1280×720 高清图像上评估），以比较模型的大小和计算复杂度。表显示，SwinIR 在不同基准数据集上，比竞争方法的 PSNR 增益最大可达到 0.53dB，同时其参数总数和乘加操作数相似。这表明 SwinIR 在图像恢复任务中具有较高的效率。

Real-world Image SR

图像超分辨率的最终目标是应用于真实世界的场景。为了测试 SwinIR 在真实世界超分辨率任务中的表现，我们使用与 BSRGAN 相同的退化模型重新训练了 SwinIR，生成低质量的图像进行合成。由于缺乏高质量的真实图像作为地面真值，我们只提供了与代表性的 双三次插值模型 ESRGAN 以及最先进的真实世界图像 SR 模型 RealSR、BSRGAN 和 Real-ESRGAN 的视觉比较。实验结果如图所示，SwinIR 生成的图像具有清晰锐利的边缘，视觉效果令人满意，而其他方法可能出现不满意的伪影。

此外，为了充分发挥 SwinIR 在真实应用中的潜力，我们进一步提出了一个更大的模型，并在更大的数据集上进行了训练。实验表明，该模型可以处理更复杂的退化问题，并在真实世界图像超分辨率任务中取得比当前模型更好的表现。

4.4 JPEG压缩伪影去除

实验设置：在 JPEG 压缩伪影去除任务中，SwinIR 与一系列基于 CNN 的方法进行了对比，实验在两个基准数据集（Classic5 和 LIVE1）上进行，测试了不同的 JPEG 压缩质量因子（10、20、30 和 40）。
结果：SwinIR 相比这些方法，特别是相比于 DRUNet（目前的最佳方法），具有更高的 PSNR（峰值信噪比）。SwinIR 只需要 11.5M 的参数，而 DRUNet 需要 32.7M 参数，显示出其在减少模型大小的同时能保持优秀性能的优势。
总结：SwinIR 在减少 JPEG 压缩伪影方面表现突出，同时模型的参数量显著较小，这表明其在压缩伪影去除上的高效性。

4.5 图像去噪

实验设置：图像去噪实验包括了对比多个方法，包括传统方法（如 BM3D 和 WNNM）和基于 CNN 的方法（如 DnCNN、IR-CNN、FFDNet 等）。在对比中使用了不同的噪声水平（15、25 和 50）。
结果：SwinIR 在去噪效果上超越了所有对比方法，尤其是在 Urban100 数据集上，超越了现有最好的模型 DRUNet 达 0.3dB。并且，SwinIR 只需 12.0M 的参数，而 DRUNet 需要 32.7M 的参数，显示了其高效的特征学习能力。
总结：SwinIR 在图像去噪方面表现优异，能够去除噪声的同时保留图像的高频细节，输出的图像边缘更加清晰，纹理更加自然。