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

来源:新智元

【导读】参数量和模型的性能有绝对关系吗?苏黎世华人博士提出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生成的图像更好,但它存在图像失真和边界伪影。

 

 

在图像降噪任务上,比较的方法包括传统模型BM3D和WNNM,基于CNN的模型DnCNN,IR-CNN,FFDNet,N3Net,NLRN,FOC-Net,RNAN,MWCNN和DRUNet。可以看到SwinIR模型比所有方法都强。

 

 

特别是它在具有100个高分辨率测试图像的大型Urban100数据集上通过最先进的DRUNet模型,最高可达0.3dB,并且SwinIR只有1200万的参数,而DRUNet有三亿参数,也就能侧面证明SwinIR的架构在学习用于图像恢复的特征表示方面是高效的。

 

 

SwinIR模型可以去除严重的噪声干扰并保留高频图像细节,从而获得更清晰的边缘和更自然的纹理。相比之下,其他方法要么过于平滑,要么过于锐利,无法恢复丰富的纹理。

参考资料:

https://arxiv.org/abs/2108.10257

推荐阅读

重磅!DLer-计算机视觉&Transformer群已成立!

大家好,这是计算机视觉&Transformer论文分享群里,群里会第一时间发布最新的Transformer前沿论文解读及交流分享会,主要设计方向有:图像分类、Transformer、目标检测、目标跟踪、点云与语义分割、GAN、超分辨率、视频超分、人脸检测与识别、动作行为与时空运动、模型压缩和量化剪枝、迁移学习、人体姿态估计等内容。

进群请备注:研究方向+学校/公司+昵称(如Transformer+上交+小明)

???? 长按识别,邀请您进群!

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值