SwinIR: Image Restoration Using Swin Transformer论文阅读

代码

Swinir: Image restoration using swin transformer

摘要：图像恢复是低水平视觉问题。目前大多图像恢复方法都是基于cnn，至于在高级视觉任务领域大放光彩的Transformer却很少使用。所以本文主要是基于Swin transformer建立了一个基准模型。这个模型包含三个部分：浅层特征提取、深层特征提取、图像重建。特别是，深度特征提取模块由几个残差Swin transformer模块（RSTB）组成，每个块都有几个Swin transformer层和一个残差连接。作者在三个代表性任务上进行了实验：图像超分辨率（包括经典、轻量级和真实世界的图像超分辨率）、图像去噪（包括灰度和彩色图像去噪）和JPEG压缩伪影减少。实验结果表明，SwinIR在不同任务上的性能优于最先进的方法，最高可达0.14∼0.45dB，而参数总数最多可减少67%。

引言

大多数基于CNN的方法侧重于精细的架构设计，例如残差学习[43，51]和密集连接[97，81]，尽管与传统的基于模型的方法[73，14，28]相比，性能有了显著提高，但它们通常面临两个基本问题，这两个问题源于基本构建块，即卷积层。首先，图像和卷积核之间的交互与内容无关。使用相同的卷积核来恢复不同的图像区域可能不是一个好的选择。

作为CNN的替代方案，Transformer[76]设计了一种自注意力机制来捕捉语境之间的全局交互，它在几个视觉问题上表现出了良好的性能[6、74、19、56]。然而，用于图像恢复的vision transformer [9，5]通常将输入图像分割成固定大小的块（例如48×48），单独处理每个块。这种策略不可避免地会产生两个缺点。首先，恢复的图像可能会在每个块周围引入边界伪影。其次，每个块的边界像素丢失信息，以便更好地恢复这可以通过重叠块来缓解，但会带来额外的计算负担。

最近，Swin Transformer[56]融合了CNN和Transformer的优点。一方面，由于局部注意力机制，它具有CNN处理大尺寸图像的优势。另一方面，它具有Transformer的优点，可以用移位窗口方案建模长期依赖。

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

基于Transformer的SwinIR具有以下优点：（1）图像内容和注意力权重之间基于内容的交互，可以解释为空间变化卷积[13，21，75]。（2） 通过移位窗口机制可以实现长期依赖性建模（3） 参数更少，性能更好。

 相关工作

图像恢复

与传统的基于模型的图像恢复方法[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]。

vision transformer

最近，自然语言处理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的图像）和多任务学习来获得良好的性能。Cao等人[5]提出了VSR Transformer，该Transformer使用自注意力机制在视频SR中进行更好的特征融合，但图像特征仍然是从CNN中提取的。此外，IPT和VSR Transformer都是局部注意，这可能不适合图像恢复。[82]提出了一种基于Swin transformer的U形架构。

方法

网络结构

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

浅层和深层特征提取

使用3*3大小的卷积层进行浅层特征提取

在早期视觉处理中卷积层很好，更稳定的优化和更好的结果。卷积可以将输入图像空间映射到高维特征空间。然后将产生的特征图输入到深层特征提取。

 其中，（·）是深度特征提取模块，它包含K个残差Swin transformer（RSTB）和一个3×3卷积层。更具体地说，中间特征F1，F2，FK和最后的输出深度特征

 其中，（·）表示第i个RSTB，是最后一个卷积层。在特征提取的尾端使用卷积层可以将卷积运算的感应偏差引入transformer，并为后面的浅层和深层特征融合奠定更好的基础。

图像重建

以图像SR为例，作者通过融合（这里是通过相加融合特征）浅层和深层特征来重建高质量的图像

 

其中，（·）是重建模块的功能。浅层特征主要包含低频信息，而深度特征侧重于恢复丢失的高频信息。通过长跳跃连接，SwinIR可以将低频信息直接传输到重建模块，这可以帮助深层特征提取模块专注于高频信息并稳定训练。对于重建模块的实现，作者使用亚像素卷积层[68]进行上采样操作。

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

 损失函数

对于图像SR，我们通过最小化L1像素损失来优化SwinIR的参数

 

 其中是通过SwinIR重建得到的，是对应的地面真实HQ图像。

对于经典和轻量级图像SR任务，只使用原始的L1像素损失函数。对于真实图像SR，我们结合使用像素损失、GAN损失和感知损失[81、89、80、27、39、81]来提高视觉质量。

对于图像去噪和JPEG压缩伪影减少，我们使用Charbonnier损失[8]。

 RSTB模块

RSTB包含swin transformer和残差块。

 

 给定第i个RSTB的输入特征，我们首先通过STL提取中间特征

 其中，（·）是第i个RSTB中的第j个STL。然后，我们在残差连接之前添加卷积层。RSTB的输出公式如下

 其中（·）是第i个RSTB中的卷积层。这种设计有两个好处。首先，虽然Transformer可以被视为空间变化卷积[21，75]的特定实例，但具有空间不变滤波器的卷积层可以增强SwinIR的平移等效性。其次，跳跃连接提供了从不同块到重建模块的快捷连接，可以融合不同层级的特征。

Swin Transformer layer 

Swin Transformer层（STL）[56]基于Transformer[76]的标准多头自注意力。和原始的transformer主要区别在于局部注意力和移动窗口。如图2（b）所示，给定大小为H×W×C的输入，Swin Transformer首先通过将输入划分为非重叠的M×M局部窗口，将输入整形为特征，其中是窗口总数。然后，它分别计算每个窗口的标准自注意力 （即局部注意力）。对于局部窗口特征，查询Q、键K和值V计算为

 其中是在不同窗口中共享的投影矩阵。一般来说，我们有。因此，注意力矩阵由局部窗口中的自注意力机制计算，如下所示

 其中B是可学习的相对位置编码。在实践中，根据[76]，作者并行执行了h次注意力函数，并连接了多头自注意力（MSA）的结果。这部分没有看懂。注意力函数和多头自注意力有什么关系。

接下来，使用具有两个完全连接层的多层感知器（MLP）进行进一步的特征变换，两个层之间具有GELU非线性激活函数。在MSA和MLP之前添加LayerNorm（LN）层，两个模块均采用残差连接。整个过程如下所述：

实验

实验设置

对于经典图像SR、真实图像SR、图像去噪和JPEG压缩伪影减少，RSTB数、STL数、窗口大小、通道数和注意头数通常分别设置为6、6、8、180和6。一个例外是，为了减少JPEG压缩伪影，窗口大小设置为7，因为我们观察到使用8时性能显著下降，可能是因为JPEG编码使用8×8图像分区。对于轻量级图像SR，我们将RSTB数和通道数分别减少到4和60。根据[95，63]，当在测试中使用self-ensemble[51]时，我们用符号“+”标记模型，例如SwinIR+。

消融实验

对于消融研究，我们在DIV2K[1]上训练SwinIR以获得经典图像SR（×2），并在Manga109[60]上进行测试。

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

我们在图中显示了通道数、RSTB数和RSTB中的STL数对模型性能的影响。据观察，峰值信噪比与这三个超参数正相关。虽然对于较大的通道数，性能不断提高，但参数总数呈二次增长。为了平衡性能和模型大小，我们在rest实验中选择180作为通道数。对于RSTB数和层数，性能增益逐渐趋于饱和。我们为它们选择了6个，以获得一个相对较小的模型。

 斑块大小和训练图像数量的影响；模型收敛性比较

我们将提出的SwinIR与一个具有代表性的基于CNN的模型RCAN进行比较，以比较基于transformer和基于CNN的差异。从图3（d）可以看出，在不同的patch大小上，SwinIR的性能优于RCAN，并且当patch较大时，峰值信噪比增益变得更大。图3（e）显示了训练图像数量的影响。当百分比大于100%（800张图像）时，将在训练中使用来自Flickr2K的额外图像。有两个观察结果。首先，正如预期的那样，SwinIR的性能随着训练图像数的增加而提高。其次，在IPT这篇论文中说基于transformer的依赖大量训练数据不同，相同训练数据SwinIR比基于CNN的模型取得了更好的结果，即使数据集很小（即25%，200幅图像）。我们还在图3（f）中绘制了SwinIR和RCAN训练期间的峰值信噪比。很明显，SwinIR比RCAN收敛更快、更好，这与之前的观察结果相矛盾，即基于变压器的模型通常存在缓慢的模型收敛

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

表1显示了RSTB中的四种残差连接变体：无残差连接、使用1×1卷积层、使用3×3卷积层和使用三个3×3卷积层（中间层的通道数设置为网络信道数的四分之一）。从表中，我们可以得到以下观察结果。首先，RSTB中的残差连接很重要，因为它将峰值信噪比提高了0.16dB。第二，使用1×1卷积带来的改善很小，可能是因为它不能像3×3卷积那样提取局部邻域信息。第三，虽然使用三个3×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上达到0.26dB。注意，RCAN和HAN引入了通道和空间注意，IGNN提出了自适应patch特征融合，NLSA基于非局部注意机制。然而，所有这些基于CNN的注意机制的性能都不如SwinIR，这表明了模型的有效性。当我们在更大的数据集上训练SwinIR时，性能进一步大幅提高（高达0.47dB），实现了比IPT更好的精度，即使IPT在训练中使用ImageNet（超过1.3M个图像），并且具有大量参数（115.5M）。相比之下，即使与最先进的基于CNN的模型（15.4~44.3）相比，SwinIR也有少量参数（11.8M）。在运行时方面，代表性的基于CNN的模型RCAN、IPT和SwinIR分别需要0.2、4.5s和1.1s在1024×1024图像上进行测试。作者在图4中显示了比例因子4的视觉对比。SwinIR可以恢复高频细节并缓解模糊，从而产生尖锐自然的边缘。相比之下，大多数基于CNN的方法无法恢复正确的纹理，并产生模糊图像或甚至不同的结构。与基于CNN的方法相比，IPT生成的图像更好，但存在图像失真和边界伪影。 

 

轻量级图像SR

作者还将SwinIR（小尺寸）与最先进的轻型图像SR方法进行了比较：CARN[2]、FALSR-A[12]、IMDN[35]、LAPAR-A[44]和LatticeNet[57]。除了峰值信噪比和SSIM，我们还报告了参数总数和乘法累加运算（在1280×720 HQ图像上评估），以比较不同模型的模型大小和计算复杂度。如表3所示，在不同的基准数据集上，参数总数和多重累积运算相似，SwinIR在峰值信噪比（PSNR）高达0.53dB，。这表明SwinIR架构是高效的。

真实世界图像SR

图像SR的最终目标是用于实际应用。最近，张等人[89]提出了一种实用的真实图像SR退化模型BSRGAN，并在真实场景中取得了非常好的结果。为了测试SwinIR对真实世界SR的性能，作者使用与BSRGAN相同的退化模型重新训练SwinIR，并在真实世界SR基准数据集RealSRSet上进行测试[89]。由于没有高质量的地面真实图像，我们仅提供与代表性双三次模型ESRGAN[81]和最先进的真实世界图像SR模型FSSR[24]、RealSR[37]和BSRGAN[89]的视觉比较。如图5所示，SwinIR生成了图像具有清晰锐利边缘，而其他比较方法可能存在不令人满意的伪影。

JPEG压缩伪影减少结果

表4显示了SwinIR与最先进的JPEG压缩伪影减少方法的比较：ARCNN[17]、DnCNN-3[90]、QGAC[20]、RNAN[96]、RDN[98]和DRUNet[88]。所有比较的方法都是基于CNN的模型。作者在两个基准数据集（Classic5[22]和LIVE1[67]）上测试。我们可以看到，对于不同的质量因子，在两个测试数据集上，提出的SwinIR的平均峰值信噪比增益至少为0.11dB和0.07dB。此外，与之前的最佳模型DRUNet相比，SwinIR只有11.5M的参数，而DRUNet是一个具有32.7M参数的大型模型。

 图像去噪结果

我们在表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。可以看出，我们的模型比所有比较的方法都具有更好的性能。特别是，在Urban100数据集上，它比最先进的模型DRUNet高达0.3dB。值得指出的是，SwinIR只有12.0M的参数，而DRUNet有32.7M的参数。这表明SwinIR结构在学习特征表示用来图像恢复方面是高效的。正如我们所见，我们的方法可以消除严重的噪声破坏，并保留高频图像细节，从而产生更清晰的边缘和更自然的纹理。相比之下，其他方法要么过于平滑，要么过于锐利，无法恢复丰富的纹理。

 总结

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

参考文献：

Transformer用于超分辨率重建_是暮涯啊的博客-CSDN博客_transformer 超分辨率