苏黎世华人博士提出模型SwinIR，只用33％的参数量就碾压图像修复领域sota

新智元报道

来源：arXiv

编辑：LRS

【新智元导读】参数量和模型的性能有绝对关系吗？苏黎世华人博士提出SwinIR模型，实验结果告诉你，越小的模型还可能更强！SwinIR使用Transformer力压CNN，又在图像修复领域屠榜，模型参数量降低67%，再也不要唯参数量论英雄了！

图像修复（image restoration）是一个受到长期关注和研究的最基础的CV问题，它能够从低质量的图像，例如缩略图、有噪音的图或是压缩图像中恢复为原始高质量图像。

但目前图像修复领域的sota方法都是基于卷积神经网络CNN的，但是很少有人尝试使用Transformer，尽管ViT在高级视觉任务中早已占据排行榜多年。

来自苏黎世联邦理工学院的华人博士提出一个适用于图像修复的模型SwinIR，主要包括浅层特征提取、深层特征提取和高质量图像重建三部分。

实验结果证明SwinIR的性能比目前sota方法提高了0.14-0.45dB，并且参数量还降低了67%。

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

项目地址：https://github.com/JingyunLiang/SwinIR

大多数基于CNN的方法侧重于精细的架构设计，如residual learning和dense learning，通过巧妙的模型设计来提升性能，增大模型容量。

虽然与传统的基于模型的方法相比CNN的性能有了显著的提高，但通常会遇到两个源于卷积层的基本问题：

1）图像和卷积核之间的交互与内容无关。使用相同的卷积核来恢复不同的图像区域可能不是最佳选择；

2）由于CNN更关注局部特征，所以卷积对于长依赖性、全局的特征建模是效果不好。

在这种情况下，很容易想到Transformer来替代CNN。Transformer的自注意力机制能够很好地捕获上下文之间的全局交互，并在多个视觉任务上具有出了良好的表现。

然而，用于图像修复的ViT需要将输入图像分割为具有固定大小（例如48×48）的patch，并对每个部分进行单独处理。

这种策略不可避免地会产生两个缺点：

1）边界像素不能利用块外的相邻像素进行图像恢复；

2）恢复的图像可能会在每个图像块周围引入边界伪影。

虽然这个问题可以通过patch重叠来缓解，但它会带来额外的计算负担。

模型设计

SwinIR的设计基于Swin Transformer，包括三个部分：

1）浅层特征抽取shallow feature extraction

浅层特征提取模块采用卷积层提取浅层特征，并将浅层特征直接传输到重构模块，以保留低频信息。

2）深层特征抽取deep feature extraction

深层特征提取模块主要由residual Swin Transformer Block（RSTB）组成组成，每个块利用多个Swin Transformer layer(STL)进行局部注意力和交叉窗口的交互。此外，还在块的末尾添加一个卷积层以增强特征，并使用残差连接为特征聚合提供快捷方式，也就是说RSTB由多个STL和一个卷积层共同构成残差块，

3）高质量图像重建high-quality（HQ） image reconstructi

重建模块是最后一步，融合了浅层和深层特征用了恢复高质量的图像。

在实验方面，作者首先研究了通道数，RSTB数目和STL数目对结果的影响。可以观察到PSNR与这三个超参数正相关。对于信道数，虽然性能不断提高，但参数量呈二次增长。为了平衡性能和模型尺寸，在剩下的实验中选择180作为通道数。对于RSTB数和层数，性能增益逐渐饱和，所以后续实验设置为6以获得一个相对较小的模型。

和经典的图像超分辨率（super-resolution, SR）模型对，包括DBPN、RCAN、RRDB、SAN、IGNN、HAN、NLSA和IPT。可以看出，当在DIV2K数据上进行训练时，SwinIR在几乎所有五个基准数据集的所有比例因子上都取得了最佳性能，在Manga109在4倍缩放上的最大PSNR增益达到0.26dB。

不过需要注意的是，RCAN和HAN引入了通道和空间注意力，IGNN提出了自适应patch特征聚合，NLSA基于非局部注意机制。所有这些基于CNN的注意机制的表现都不如所提出的基于Transformer的SwinIR，这也表明了文中所提出模型的有效性。

当在更大的数据集（DIV2K+Flickr2K）上训练SwinIR时，性能进一步大幅度提高，也实现了比基于Transformer的模型IPT更好的精度，达到0.47dB。即使IPT在训练中使用ImageNet（超过一百三十万图像），并且具有超过一亿的参数。相比之下，即使与基于CNN的sota模型相比，SwinIR的参数也很少（1500万-4430万）。

在运行时方面，与基于CNN的代表性的模型RCAN相比，IPT和SwinIR在1024×1024分别需要约0.2、4.5和1.1秒。

实验结果

从可视化结果来看，SwinIR可以恢复高频细节并减轻模糊瑕疵，并产生锐利且自然的边缘。

相比之下，大多数基于CNN的方法会产生模糊的图像，甚至是不正确的纹理。与基于CNN的方法相比，IPT生成的图像更好，但它存在图像失真和边界伪影。