• 问答
  • 技术
  • 实践
  • 资源
吊打 IPT,ETH 学者将 SwinTransformer 应用图像超分,图像降噪、JPEG 压缩伪影移除,性能大幅提升
技术讨论

作者丨happy
编辑丨极市平台

原文链接: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在真实场景图像超分方面的应用与效果,期待作者尽快开源。

Abstract

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

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

相比CNN方案,Swin具有以下几个优势:

  • 基于内容交互的图像内容与注意力权值可以视作空域可变卷积;
  • RSTB中的移位窗口机制可以进行长距离依赖建模;
  • 更优的性能、更少的参数(可参见下图:SwinIR具有更少的参数量、更优的性能)。

image.png

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

Method

Network Architecture

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

Shallow and deep feature extraction. 给定低质输入$I{LQ} \in R^{H \times W \times C{in}}$,我们采用$3\times 3$卷积$H_{SF}(\cdot)$提取浅层特征$F_0 \in R^{H \times W \times C}$:

$F0 = H{SF}(I_{LQ})$

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

$F{DF} = H{DF}(F_0)$

注:$H_{DF}$表示深层特征提取模块,它由K个RSTB与一个$3\times 3$卷积构成。更具体来说,中间特征$F_1, F_2, \cdots, FK$以及深层特征$F{DF}$的计算可以描述如下:

$Fi = H{RSTBi}(F{i-1}), i=1,2,\cdots, K \F{DF} = H{CONV}(F_K)$

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

Image Reconstruction. 以图像超分为例,我们通过聚合浅层与深层特征进行高质量图像$I_{RHQ}$重建:

$I{RHQ} = H{REC}(F0 + F{DF})$

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

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

$I{RHQ} = H{SwinIR}(I{LQ}) + I{LQ}$

Loss function. 对于图像超分,我们采用$L_1$损失进行优化:

$\mathcal{L} = | I{RHQ} - i{HQ} |_1$

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

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

Residual Swin Transformer Block

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

$F{i,j} = H{STL{i,j}}(F{i,j-1}), j=1,2,\cdots, L$

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

$F{i,out} = H{CONVi}(F{i,L}) + F_{i,0}$

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

  • 尽管Transformer可以视作空间可变卷积的变种,但空间不变卷积有助于提升SwinIR的平移不变形;
  • 残差连接为不同模块到重建模块提供了等效连接,促进了不同层级特征的聚合。

Swin Transformer layer. Swin Transformer Layer(STL)基于原始Transformer中的标准多头自注意力演变而来,主要区别体现在于局部注意力移位窗口机制。前述图b给出了STL结构示意图。假设输入尺寸为$H\times W\times C$,首先,将输入拆分为$M\times M$局部窗口并reshape为$\frac{HW}{M^2} \times M^2 \times C$;然后,在每个窗口计算标准自注意力。对于局部窗口特征$X \in R^{M^2 \times C}$,query、key以及value计算如下:

$Q = XP_Q, K=XP_K, V = XP_V$

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

$Attention(Q,K,V) = SoftMax(QK^T/\sqrt{d} + B)V$

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

$X = MSA(LN(X)) + X \X = MLP(LN(X)) + X$

然而,当不同层的窗口划分固定时,不同局部窗口之间不存在新交互。因此,窗口划分与窗口移位用于进行跨窗口信息交互。更多关于SwinTransformer的介绍可参见【zzk:图解Swin Transformer】

Experiments

Main Results

Classical Image SR. 上表给出了经典图像超分方面的性能对比,从中可以看到:

  • 当在DIV2K数据集上训练时,在五个基准数据集上,Swin取得了所有尺度最佳性能;PSNR增益最高甚至可达0.26dB。
  • 当在DIV2K+Flickr2K数据集上训练时,SwinIR的性能得到了进一步提升,超越了ImageNet数据集上训练的IPT;
  • 值得注意的是,SwinIR的参数量远小于IPT(11.8M vs 115.5M),甚至比CNN方案的参数量(15.4~44.3M)还少
  • 在推理速度方面,当输入为$1024 \times 1024$时,RCAN、IPT以及SwinIR速度分别为0.2s、4.5s以及1.1s;

image.png

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

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

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

JPEG Compression Artifact Reduction. 上表给出了JPEG压缩伪影移除任务上的性能对比,从中可以看到:

  • 在两个数据集上,所提SwinIR分别取得了平均0.11dB、0.07dB指标提升;
  • 相比此前最佳DRUNet(参数量为32.7M),SwinIR参数量仅为11.5M参数。

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

  • 相比其他方案,所提SwinIR在图像降噪任务上同样具有更优异的性能;在Urban100数据上,SwinIR超出DRUNet0.3dB;
  • 在参数量方面,SwinIR仅需12M,而DRUNet则具有32.7M;这说明:SwinIR在学习用于复原的特征表达上非常高效

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

Ablation Study

Impact of channel number, RSTB number and STL number.

上图abc对通道数、模块数以及STL数进行消融分析对比,从中可以看到:

  • PSNR指标与上述三个超参成正相关关系;
  • 对于通道数而言,尽管通道数越多性能越好,但总参数成二次方关系增长。为平衡性能与模型大小,我们设置通道数180;
  • 随RSTB模块数与STL数增加,模型性能很快饱和。我们将两者数量设为6以得到一个相对小的模型。

Impact of patch size and training image number; model convergence comparson.

上图def对训练过程中的超参进行了消融分析对比,从中可以看到:

  • 在不同patch尺寸下,SwinIR均取得了比RCAN更佳的性能,见上图d;
  • SwinIR性能会随训练数据量提升而提升且均比RCAN更优,见上图e;
  • SwinIR的收敛速度要比RCAN更快、更好,见上图f。

Impact of residual connection and convolution layer in RSTB.

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

  • 残差连接非常重要,它可以提升的性能高达0.16dB;
  • $1\times 1$卷积带来的性能增强非常有限,这是因为它不会像$3\times 3$卷积一样提取局部近邻信息;
  • 尽管采用三个$3\times 3$卷积(减少中间卷积通道数)可以降低参数量,但性能同样轻微下降。
  • 0
  • 0
  • 732
收藏
暂无评论
sophie
大咖

科技园的搬砖汪

  • 18,113

    关注
  • 328

    获赞
  • 54

    精选文章
近期动态
  • 从事AI视觉算法开发多年,主攻目标检测、图像分割方向
文章专栏
  • 优质论文推荐