SwinIR: Image Restoration Using Swin Transformer

Abstract

图像复原是一个由来已久的低层视觉问题，其目标是从低质量图像（例如降采样、带噪声或压缩图像）中还原出高质量图像。尽管当前最先进的图像复原方法大多基于卷积神经网络，但在高层视觉任务中表现出色的Transformer尚少有尝试应用于此领域。

在本文中，我们提出了一个基于Swin Transformer的强基线模型——SwinIR，用于图像复原任务。SwinIR主要由三个部分组成：浅层特征提取、深层特征提取和高质量图像重建。其中，深层特征提取模块由多个残差Swin Transformer块（Residual Swin Transformer Block，RSTB）构成，每个块内部包含多个Swin Transformer层，并配有残差连接结构。

我们在三个具有代表性的任务上进行了实验：图像超分辨率（包括经典、轻量级和真实场景图像超分辨率）、图像去噪（包括灰度图像和彩色图像去噪）以及JPEG压缩伪影去除。实验结果表明，SwinIR在多个任务中均优于现有最先进的方法，性能提升幅度可达0.14∼0.45dB，同时模型参数总量最多可减少67%。

1. Introduction

图像复原任务（如图像超分辨率（SR）、图像去噪和JPEG压缩伪影去除）旨在从低质量退化图像中重建出高质量的干净图像。自从一系列具有革命性影响的工作问世以来 [18, 40, 90, 91]，卷积神经网络（CNN）已经成为图像复原的主力方法 [43, 51, 43, 81, 92, 95, 24, 93, 46, 89, 88]。大多数基于CNN的方法主要聚焦于精心设计的网络架构，如残差学习 [43, 51] 和密集连接 [97, 81]。虽然与传统的基于模型的方法相比，这些方法在性能上有了显著提升 [73, 14, 28]，但它们普遍存在两个源自卷积层本身的基本问题：

1. 图像与卷积核之间的交互是内容无关的。对不同图像区域使用相同的卷积核可能并非最优选择。
2. 由于卷积遵循局部处理原则，它在建模长距离依赖时效果有限。

作为CNN的替代方案，Transformer [76] 设计了一种自注意力机制，用于捕捉全局上下文之间的交互，并在多个视觉任务中展现出良好的性能 [6, 74, 19, 56]。然而，应用于图像复原的视觉Transformer方法 [9, 5] 通常将输入图像划分为固定大小的图像块（如48×48），并独立处理每个图块。这种策略不可避免地带来两个缺陷：

1. 图块边缘像素无法利用图块以外的邻近像素进行图像复原；
2. 重建后的图像可能在每个图块边缘出现伪影。

虽然通过图块重叠可以缓解这一问题，但这会显著增加计算负担。

近年来，Swin Transformer [56] 展现出极大的潜力，它结合了CNN与Transformer的优点：

• 一方面，借助局部注意力机制，它具备CNN在处理大尺寸图像方面的优势；
• 另一方面，通过滑动窗口机制（shifted window scheme），它能够有效建模长距离依赖。

在本文中，我们基于Swin Transformer提出了一种图像复原模型，称为SwinIR。具体来说，SwinIR包含三个模块：

1. 浅层特征提取模块：使用一个卷积层提取浅层特征，并将其直接传递给重建模块，以保留低频信息；
2. 深层特征提取模块：主要由多个残差Swin Transformer块（RSTB）组成，每个块包含多个Swin Transformer层，用于局部注意力和跨窗口交互。此外，我们在每个块末尾加入一个卷积层以增强特征，并使用残差连接以实现特征聚合的捷径；
3. 高质量图像重建模块：将浅层与深层特征融合，用于最终的高质量图像重建。

与主流的CNN图像复原模型相比，基于Transformer的SwinIR具有以下几个优势：

1. 基于内容的交互机制：图像内容与注意力权重的交互可以被解释为一种空间变化的卷积 [13, 21, 75]；
2. 长距离依赖建模能力：得益于滑动窗口机制；
3. 在更少参数下实现更优性能：例如，如图1所示，SwinIR在参数更少的情况下，相比现有图像超分辨率方法，能获得更高的PSNR（峰值信噪比）性能。

2. Related Work

2.1. Image Restoration

与传统的图像复原方法 [28, 72, 73, 62, 32]（通常是基于模型的方法）相比，基于学习的方法，尤其是基于CNN的方法，由于其出色的性能，已变得更加流行。这类方法通常通过大规模成对的数据集，学习低质量图像与高质量图像之间的映射关系。

自从开创性的工作——SRCNN[18]（用于图像超分辨率）、DnCNN [90]（用于图像去噪）以及ARCNN [17]（用于JPEG压缩伪影去除）问世以来，涌现出大量基于CNN的模型，旨在通过更复杂的神经网络结构设计提升模型的表示能力。

这些结构设计包括：

• 残差块（Residual Block）[40, 7, 88]；
• 密集块（Dense Block）[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框架中引入了注意力机制，例如：

• 通道注意力（Channel Attention）[95, 15, 63]；
• 非局部注意力（Non-local Attention）[52, 61]；
• 自适应图块聚合（Adaptive Patch Aggregation）[100]。

2.2 Vision Transformer

近年来，自然语言处理中的模型 Transformer [76] 在计算机视觉领域变得越来越受欢迎。当应用于图像分类 [66, 19, 84, 56, 45, 55, 75]、目标检测 [6, 53, 74, 56]、图像分割 [84, 99, 56, 4] 和人群计数 [47, 69] 等视觉任务时，Transformer 通过探索图像不同区域之间的全局交互，学习关注关键图像区域。得益于其出色的表现，Transformer 也被引入到图像复原任务中 [9, 5, 82]。Chen 等人 [9] 基于标准 Transformer 提出了一个主干模型 IPT，用于解决多种图像复原问题。然而，IPT 的性能依赖于大量参数（超过 1.155 亿个参数）、大规模数据集（超过 110 万张图像）以及多任务学习策略。Cao 等人 [5] 提出了 VSR-Transformer，该模型在视频超分辨率任务中利用自注意力机制实现更优的特征融合，但其图像特征仍然通过 CNN 提取。此外，IPT 和 VSR-Transformer 都采用基于图像块的注意力机制，这可能并不适用于图像复原任务。另外，一项同期的工作 [82] 提出了一种基于 Swin Transformer [56] 的 U 形结构。

3. Method

3.1 Network Architechture

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

浅层与深层特征提取

给定一个低质量（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 表示特征通道数。该卷积层在图像处理初期表现良好，有助于更稳定的优化和取得更优的结果 [86]。它还提供了一种简单的方法，将输入图像空间映射到更高维的特征空间。

接着，我们从 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, \ldots, F_K 以及最终的深层特征 F_{DF} 是逐块提取的，计算如下：

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

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

其中 H_{RSTB_i}(\cdot) 表示第 i 个 RSTB（残差 Swin Transformer 块），H_{CONV} 是最后的卷积层。在特征提取末尾使用卷积层，可以将卷积操作的归纳偏置引入到基于 Transformer 的网络中，为后续的浅层与深层特征融合打下更好的基础。

Image reconstruction

以图像超分辨（SR）为例，我们通过聚合浅层和深层特征来重建高质量图像 I_{RHQ}，计算方式如下：

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

其中 H_{REC}(\cdot) 是重建模块的函数。浅层特征主要包含低频信息，而深层特征则专注于恢复丢失的高频信息。借助长跳跃连接，SwinIR 可以将低频信息直接传递给重建模块，有助于深层特征提取模块专注于高频信息并稳定训练。

在重建模块的实现中，我们使用亚像素卷积层（sub-pixel convolution layer）[68] 来对特征进行上采样。

对于不需要上采样的任务，如图像去噪和 JPEG 压缩伪影去除，仅使用一个卷积层进行重建。此外，我们采用残差学习来重建 LQ 图像与 HQ 图像之间的残差，而不是直接重建 HQ 图像。其公式如下：

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

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

Loss function

对于图像超分任务，我们通过最小化 L_1 像素损失来优化 SwinIR 的参数：

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

其中 I_{RHQ} 是以 I_{LQ} 为输入，通过 SwinIR 得到的结果，I_{HQ} 是对应的真实高质量图像。

对于经典或轻量级图像超分任务，我们仅使用简单的 L_1 像素损失与以往工作保持一致，以展示所提出网络的有效性。对于真实图像超分任务，我们结合像素损失、GAN 损失以及感知损失 [81, 89, 80, 27, 39, 81] 来提升视觉质量。

对于图像去噪和 JPEG 压缩伪影去除任务，我们使用 Charbonnier 损失 [8]，其定义如下：

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

其中 \epsilon 是一个常数，经验上设为 10^{-3}。

3.2. Residual Swin Transformer Block

如图 2(a) 所示，残差 Swin Transformer 块（RSTB）是一个包含 Swin Transformer 层（STL）和卷积层的残差结构。给定第 i 个 RSTB 的输入特征 F_{i,0}，我们首先通过 L 个 Swin Transformer 层提取中间特征 F_{i,1}, F_{i,2}, \dots, F_{i,L}，计算如下：

F_{i,j} = H_{STL_{i,j}}(F_{i,j-1}), \quad j = 1, 2, \dots, L

其中 H_{STL_{i,j}}(\cdot) 是第 i 个 RSTB 中第 j 个 Swin Transformer 层。接着，在残差连接前添加一个卷积层，RSTB 的输出定义为：

F_{i,\text{out}} = H_{CONV_i}(F_{i,L}) + F_{i,0}

其中 H_{CONV_i}(\cdot) 是第 i 个 RSTB 中的卷积层。这一设计有两个好处：

1. 尽管 Transformer 可以看作是空间变化卷积的一种特例 [21, 75]，但卷积层具有空间不变性的滤波器，可以增强 SwinIR 的平移等变性；
2. 残差连接提供了一种基于恒等映射的连接方式，将不同模块的信息传递到重建模块，有助于融合不同层级的特征。

Swin Transformer layer

Swin Transformer 层（STL）[56] 基于原始 Transformer 层的标准多头自注意力机制（multi-head self-attention）[76]。其主要区别在于局部注意力和滑动窗口机制。

如图 2(b) 所示，给定输入尺寸为 H \times W \times C，Swin Transformer 首先将输入划分为不重叠的 M \times M 局部窗口，得到大小为 \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、K 和 V 的计算如下：

Q = XP_Q, \quad K = XP_K, \quad V = XP_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{QK^T}{\sqrt{d}} + B\right)V

其中 B 是可学习的相对位置编码。在实际操作中，遵循 [76] 的做法，我们并行执行 h 次注意力计算，并将结果拼接用于多头自注意力（MSA）。

接下来是一个多层感知机（MLP），它包含两个带有 GELU 非线性激活的全连接层，用于进一步的特征变换。在 MSA 和 MLP 之前都添加了 LayerNorm（LN）层，且两个模块都使用了残差连接，整个过程如下：

X = \text{MSA}(\text{LN}(X)) + X,

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

然而，当每层使用固定划分方式时，不同窗口之间无法建立联系。因此，交替使用规则划分和滑动窗口划分来实现跨窗口连接 [56]。滑动窗口划分的含义是：在划分前将特征图沿两个方向分别平移 \left(\left\lfloor \frac{M}{2} \right\rfloor, \left\lfloor \frac{M}{2} \right\rfloor\right) 个像素。

4. Experiment

4.1. Experimental Setup

对于经典图像超分、真实图像超分、图像去噪和 JPEG 压缩伪影去除等任务，RSTB 数量、STL 数量、窗口大小、通道数和注意力头数通常设置为 6、6、8、180 和 6。其中一个例外是：在 JPEG 压缩伪影去除任务中，窗口大小设置为 7，因为我们在使用 8 时性能显著下降，可能是因为 JPEG 编码采用的是 8 \times 8 的图像分块方式。

对于轻量级图像超分任务，我们将 RSTB 数量和通道数分别减少为 4 和 60。参考 [95, 63]，在测试过程中使用自集成策略（self-ensemble strategy）[51] 时，我们使用符号 “+” 表示模型，例如 SwinIR+。由于篇幅限制，训练与评估细节在补充材料中提供。

4.2. Ablation Study and Discussion

为了进行消融实验，我们在 DIV2K 数据集上训练 SwinIR（×2 放大）并在 Manga109 上进行测试 [1, 60]。

Impact of channel number, RSTB number and STL number

我们在 RSTB 中研究了通道数、RSTB 数量和 Swin Transformer 层数对模型性能的影响，结果如图 3(a)、3(b)、3(c)。实验表明，PSNR 与这三个超参数呈正相关。

• 通道数：随着通道数增加，性能持续提升，但参数量呈平方增长。为了在性能与模型规模之间取得平衡，我们在其余实验中将通道数设为 180。
• RSTB 数量与层数：性能提升逐渐趋于饱和，最终选择二者均为 6，以构建相对较小的模型。

Impact of patch size and training image number; model convergence comparison

我们将 SwinIR 与典型 CNN 模型 RCAN 进行对比，研究 Transformer 与 CNN 模型的差异。

• 从图 3(d) 可看出，SwinIR 在不同 patch size 设置下表现均优于 RCAN，patch 越大，PSNR 提升越明显。
• 图 3(e) 显示训练图像数量的影响，使用额外的 Flickr2K 图像进行训练。结果表明：
  • SwinIR 性能随着训练图像数量增加而提升；
  • 与 IPT 中 Transformer 模型对大量数据的依赖不同，SwinIR 在训练数据较少（如 25%，200 张图像）时也优于 CNN 模型；
• 图 3(f) 展示了 SwinIR 与 RCAN 的训练过程中 PSNR 变化曲线，显示 SwinIR 收敛更快，性能更好，打破了 Transformer 模型收敛慢的常规认知。

Impact of residual connection and convolution layer in RSTB

表 1 展示了 RSTB 中四种不同的残差连接变体：

• 无残差连接；
• 使用 1 \times 1 卷积；
• 使用 3 \times 3 卷积；
• 使用三个 3 \times 3 卷积（中间通道为总通道数的四分之一）。

结论如下：

1. 残差连接对性能有益，可提升 PSNR 约 0.16dB；
2. 1 \times 1 卷积提升有限，可能是因为无法像 3 \times 3 卷积那样提取邻域信息；
3. 使用三个 3 \times 3 卷积可减少参数数目，但性能略有下降。

4.3. Results on Image SR

Classical image SR

表 2 显示了 SwinIR（中等规模）与当前最优方法的定量对比，包括 DBPN [31]、RCAN [95]、RRDB [81]、SAN [15]、IGNN [100]、HAN [63]、NLSA [61] 和 IPT [9]。

• 在 DIV2K 上训练后，SwinIR 在几乎所有基准数据集和尺度上都取得最佳性能；
• 在 Manga109 上，放大因子为 4 的 PSNR 最大提升达 0.26dB；
• 尽管 RCAN 和 HAN 引入了通道与空间注意力机制，IGNN 使用了自适应 patch 特征聚合，NLSA 采用非局部注意力，但它们都不如基于 Transformer 的 SwinIR；
• 使用更大数据集（DIV2K+Flickr2K）训练后，SwinIR 性能进一步提升，最高提升 0.47dB，超过了使用 ImageNet（130 万图像）训练并拥有 115.5M 参数的 IPT 模型；
• 相比之下，SwinIR 仅有 11.8M 参数，比 CNN 模型（15.4~44.3M）更小；
• 在测试时间上，RCAN、IPT 和 SwinIR 分别为 0.2s、4.5s 和 1.1s（在 1024 \times 1024 图像上），如图 4 所示，SwinIR 能恢复高频细节并减少模糊伪影，生成清晰自然的图像。

Lightweight image SR

我们还将 SwinIR（小规模）与当前轻量级 SR 方法进行对比：CARN [2]、FALSR-A [12]、IMDN [35]、LAPAR-A [44] 和 LatticeNet [57]。

• 除了 PSNR 和 SSIM，我们还报告了参数数量和乘加运算次数（在 1280 \times 720 HQ 图像上评估）；
• 如表 3 所示，SwinIR 在多个基准数据集上最多提高 0.53dB，且参数量与计算量相近；
• 说明 SwinIR 架构在图像复原任务中具有极高效率。

Real-world image SR

图像超分的最终目标是实际应用。近期，Zhang 等人 [89] 提出实用的退化模型 BSRGAN，并在真实图像 SR 任务中表现优异。

• 为评估 SwinIR 在真实图像 SR 的性能，我们使用与 BSRGAN 相同的退化模型对 SwinIR 进行再训练；
• 由于没有真实高质量图像，我们仅与代表性模型（如 ESRGAN [81]、RealSR [37]、BSRGAN [89]、Real-ESRGAN [80]）进行视觉对比；
• 如图 5 所示，SwinIR 生成的图像具有清晰锐利的边缘和良好的感官质量，而其他方法可能存在伪影；
• 为进一步发掘 SwinIR 在真实场景中的潜力，我们训练了一个更大的模型，并在更大数据集上进行训练。

实验表明，该大模型能够处理更复杂的退化问题，在真实图像上表现更佳。详细内容请参见项目主页： https://github.com/JingyunLiang/SwinIR

4.4. Results on JPEG Compression Artifact Reduction

表 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。此外，SwinIR 仅有 11.5M 参数，而此前最佳模型 DRUNet 拥有 32.7M 参数，显著更轻量。

4.5. Results on Image Denoising

我们在表 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 的情况下进行对比。

实验表明，SwinIR 在所有对比方法中表现最好，尤其是在包含 100 张高分辨率图像的 Urban100 数据集上，PSNR 提升高达 0.3dB，超过当前最优模型 DRUNet。

值得注意的是，SwinIR 仅有 12.0M 参数，而 DRUNet 有 32.7M 参数，这表明 SwinIR 在学习用于图像复原的特征表示方面具有高度效率。

灰度和彩色图像的可视化对比结果如图 6 和图 7 所示。实验表明，SwinIR 能够有效去除严重噪声扰动，保留图像中的高频细节，从而获得更清晰的边缘和自然的纹理。而其他方法要么过于平滑，要么锐化过度，无法恢复丰富的纹理信息。

5. Conclusion

在本文中，我们提出了一种基于 Swin Transformer 的图像复原模型——SwinIR。该模型由三个部分组成：浅层特征提取、深层特征提取和高质量图像重建模块。

特别地，我们为深层特征提取使用了一组残差 Swin Transformer 块（RSTB），每个 RSTB 由 Swin Transformer 层、卷积层和残差连接组成。

大量实验证明，SwinIR 在三类具有代表性的图像复原任务中，在六种不同设置下均达到了当前最先进的性能。这些设置包括：经典图像超分辨（SR）、轻量级图像 SR、真实图像 SR、灰度图像去噪、彩色图像去噪和 JPEG 压缩伪影去除。这充分展示了所提出的 SwinIR 的有效性和泛化能力。

未来，我们将扩展该模型至其他图像复原任务，例如图像去模糊和去雨等。