Transformer在图像复原领域的又一力作！ETH提出SwinIR：low-level视觉多项任务全面领先...

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

作者丨happy

转载自丨极市平台

导读

 

鲜少有工作将transformer用于图像复原方向，之前有篇：Transformer再下一城！low-level多个任务榜首被占领，中科大等联合提出：Uformer

近期ETH的学者近日提出基于Swin Transformer的一种强基线模型SwinIR用于图像复原，该工作在经典图像超分、真实场景图像超分、图像降噪与JPEG压缩伪影移除都取得了显著优于已有方案的性能。

SwinIR: Image Restoration Using Swin Transformer

原文链接：https://arxiv.org/abs/2108.10257

code：https://github.com/JingyunLiang/SwinIR（未开源）

期待已久的SwinIR终于刊出来了，它是ETH团队在Transformer+low-level方面的最新力作，大幅超越了之前华为诺亚提出的IPT。从某种程度上来看，这篇文章没有什么创新点，只是将微软的Swin Transformer在low-level任务上进行了应用。Whatever，无论是经典图像超分(即退化方式为bicubic)，还是真实场景图像超分，亦或图像降噪与JPEG压缩伪影移除，所提SwinIR均取得了显著优于已有方案的性能。笔者这里尤其期待SwinIR在真实场景图像超分方面的应用与效果，期待作者尽快开源。

摘要

图像复原(存在已久的low-level视觉问题)旨在根据低质图像(比如，下采样的、带噪的、压缩的图像)复原高质量图像。现有优异图像复原方案往往采用CNN，鲜少有Transformer(已在high-level视觉任务中取得骄人成绩)类方案在该类问题中进行探索尝试。

本文基于Swin Transformer提出一种强基线模型SwinIR用于图像复原。SwinIR包含三部分：浅层特征提取、深层特征提取以及高质量图像重建。具体而言，深层特征提取由多个RSTB(Residual Swin Transformer Blocks)构成，每个RSTB由多个Swin Transformer层与残差连接构成。

相比CNN方案，Swin具有以下几个优势：

• 基于内容交互的图像内容与注意力权值可以视作空域可变卷积；

• RSTB中的移位窗口机制可以进行长距离依赖建模；

• 更优的性能、更少的参数(可参见下图：SwinIR具有更少的参数量、更优的性能)。

我们在三个极具代表性的任务(图像超分、图像降噪以及JPEG压缩伪影移除)上进行了实验。实验结果表明：所提SwinIR能够以0.14~0.45dB优于其他SOTA方案，同时参数量降低高达67%。

方法简介

网络结构

上图给出了本文所提SwinIR网络架构示意图，很明显，它由三个部分构成：浅层特征提取、深层特征提取以及高质量图像重建模块。对于不同任务而言，所提方案的区别主要在于重建模块。

浅层及深层特征提取

给定低质输入 ，我们采用 卷积 提取浅层特征 :

已有研究[ConvStem]表明: 卷积善于进行早期视觉处理，同时有助于稳化训练并取得更优结果。此外，卷积还提供了一种将输入图像映射到更高维特征空间的简单方案。然后，我们对 提取深层特征 ：

注： 表示深层特征提取模块，它由K个RSTB与一个 卷积构成。更具体来说，中间特征 以及深层特征 的计算可以描述如下：

注：这里的卷积操作可以将卷积操作的归纳偏置引入到SwinIR中，同时为浅层特征与深层特征的聚合奠定更好的基础。

图像重建

以图像超分为例，我们通过聚合浅层与深层特征进行高质量图像 重建:

注： 表示重建模块。由于浅层特征主要包含低频信息，而深层特征聚焦于重建遗失的高频信息，SwinIR采用了长距离跳过连接将两者聚合进行最终的重建。

对于超分而言，重建模块由卷积与PixelShuffle构成；对于不需要上采样的任务(如降噪、JPEG压缩伪影移除)，重建模块仅由卷积构成。此外，我们采用残差学习机制重建LQ与HR之间的残差，即：

损失函数

对于图像超分，我们采用 损失进行优化：

对于经典与轻量图像超分，我们近采用上述 损失；对于真实世界图像超分，我们采用 损失、GAN损失以及感知损失的组合以提升视觉质量。对于图像降噪与JPEG压缩伪影任务，我们采用Charbonnier损失：

残差Swin Transformer模块

上图a给出了RSTB的结构示意图，它包含多个STL、一个卷积以及残差连接。对于第 个RSTB，输入特征表示为 ，我们首先通过L个Swin Transformer层提取中间特征 :

然后，我们在残差连接之前添加一个卷积层。RSTB的输出可以描述如下：

这种设计思路有这样两个优势：

1) 尽管Transformer可以视作空间可变卷积的变种，但空间不变卷积有助于提升SwinIR的平移不变形；

2) 残差连接为不同模块到重建模块提供了等效连接，促进了不同层级特征的聚合。

Swin Transformer层

Swin Transformer Layer(STL)基于原始Transformer中的标准多头自注意力演变而来，主要区别体现在于局部注意力与移位窗口机制。前述图b给出了STL结构示意图。假设输入尺寸为 ，首先，将输入拆分为 局部窗口并reshape为 ；然后，在每个窗口计算标准自注意力。对于局部窗口特征 ，query、key以及value计算如下：

基于上述信息，注意力矩阵的计算公式如下：

其中，B表示可学习相对位置编码。接下来，我们采用包含两个全连接层与GELU激活的MLP进行特征变换。因此，整个过程可以描述如下：

然而，当不同层的窗口划分固定时，不同局部窗口之间不存在新交互。因此，窗口划分与窗口移位用于进行跨窗口信息交互。更多关于SwinTransformer的介绍

可参见：zzk：图解Swin Transformer，https://zhuanlan.zhihu.com/p/367111046

实验

主要结果

经典图像超分

上表给出了经典图像超分方面的性能对比，从中可以看到：

1) 当在DIV2K数据集上训练时，在五个基准数据集上，Swin取得了所有尺度最佳性能；PSNR增益最高甚至可达0.26dB。

2) 当在DIV2K+Flickr2K数据集上训练时，SwinIR的性能得到了进一步提升，超越了ImageNet数据集上训练的IPT;

3) 值得注意的是，SwinIR的参数量远小于IPT(11.8M vs 115.5M)，甚至比CNN方案的参数量（15.4~44.3M）还少；

4) 在推理速度方面，当输入为 时，RCAN、IPT以及SwinIR速度分别为0.2s、4.5s以及1.1s;

上图给出了SwinIR与其他超分方案的视觉效果对比，可以看到：SwinIR可以复原高频细节，具有更锐利而自然的边缘；相反，CNN方案生成结果往往具有模糊结果，甚至不正确的纹理。

轻量图像超分

上图对比了不同轻量型方案的性能对比，从中可以看到：SwinIR显著优于其他方案(PSNR指标最高超出0.53dB)，同时具有相近的参数量与计算量。

真实世界图像超分

图像超分的终极目标是真实场景实用。我们以BSRGAN为基础进行SwinIR训练，视觉效果对比见上图。从图示效果可以看到：SwinIR生成了视觉效果更佳(更清晰、边缘更锐利)的图像，而其他方案则存在伪影问题。对标的模型包含：

• ESRGAN；更详细信息可参见【科技猛兽：ESRGAN原理分析和代码解读】(https://zhuanlan.zhihu.com/p/156505590)；

• BSRGAN：更详细信息可参见【真正实用的退化模型，ETH开源业界首个广义盲图像超分退化模型】；

• RealSR：更细信息可参见【显著提升真实数据超分性能，南大&腾讯开源图像超分新方案RealSR】；

• Real-ESRGAN：更详细可参见【Real-ESRGAN：走出温室，迈向实用！】

JPEG压缩伪影移除

上表给出了JPEG压缩伪影移除任务上的性能对比，从中可以看到：

• 在两个数据集上，所提SwinIR分别取得了平均0.11dB、0.07dB指标提升；

• 相比此前最佳DRUNet(参数量为32.7M)，SwinIR参数量仅为11.5M参数。

图像降噪

上面两个表给出了灰度/彩色图像降噪方面的性能对比，从中可以看到：

• 相比其他方案，所提SwinIR在图像降噪任务上同样具有更优异的性能；在Urban100数据上，SwinIR超出DRUNet0.3dB；

• 在参数量方面，SwinIR仅需12M，而DRUNet则具有32.7M；这说明：SwinIR在学习用于复原的特征表达上非常高效。

上面两个图给出了灰度/彩色图像降噪方面的视觉效果对比，从中可以看到：SwinIR可以有效移除重度噪声干扰，同时保留高频图像细节，进而生成更锐利的边缘、更自然的纹理；而其他方案要么过于平滑，要么过度锐化，难以重建丰富纹理。

消融实验

对通道数量、RSTB 模块数和 STL 数的影响

上图a、b、c对通道数、模块数以及STL数进行消融分析对比，从中可以看到：

• PSNR指标与上述三个超参成正相关关系；

• 对于通道数而言，尽管通道数越多性能越好，但总参数成二次方关系增长。为平衡性能与模型大小，我们设置通道数180；

• 随RSTB模块数与STL数增加，模型性能很快饱和。我们将两者数量设为6以得到一个相对小的模型。

上图d、e、f对训练过程中的超参进行了消融分析对比，从中可以看到：

• 在不同patch尺寸下，SwinIR均取得了比RCAN更佳的性能，见上图d；

• SwinIR性能会随训练数据量提升而提升且均比RCAN更优，见上图e；

• SwinIR的收敛速度要比RCAN更快、更好，见上图f。

上表对RSTB的各成分进行了消融对比，从中可以看到：

• 残差连接非常重要，它可以提升的性能高达0.16dB；

• 卷积带来的性能增强非常有限，这是因为它不会像 卷积一样提取局部近邻信息；

• 尽管采用三个 卷积(减少中间卷积通道数)可以降低参数量，但性能同样轻微下降。

本文亮点总结

1.SwinIR包含三部分：浅层特征提取、深层特征提取以及高质量图像重建。具体而言，深层特征提取由多个RSTB(Residual Swin Transformer Blocks)构成，每个RSTB由多个Swin Transformer层与残差连接构成。

2.相比CNN方案，Swin具有以下几个优势：

• 基于内容交互的图像内容与注意力权值可以视作空域可变卷积；

• RSTB中的移位窗口机制可以进行长距离依赖建模；

• 更优的性能、更少的参数

SwinIR论文下载

后台回复：SwinIR，即可下载上述论文PDF

后台回复：ICCV2021，即可下载ICCV 2021论文和代码开源的论文合集

后台回复：CVPR2021，即可下载CVPR 2021论文和代码开源的论文合集

CVer-Transformer交流群成立

扫码添加CVer助手，可申请加入CVer-Transformer 微信交流群，方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch和TensorFlow等群。

一定要备注：研究方向+地点+学校/公司+昵称（如Transformer+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲长按加小助手微信，进交流群

▲点击上方卡片，关注CVer公众号

整理不易，请点赞和在看