NTIRE2020 冠军方案 RFB-ESRGAN:带感受野模块的超分网络

Tips
◎作者系极市原创作者计划特约作者Happy
欢迎大家联系极市小编(微信ID:fengcall19)加入极市原创作者行列

paper: https://arxiv.org/abs/2005.12597

本文是OPPO研究院提出并参加NTIRE2020Perceptual Extreme Super Resolution(一般泛指x16超分)竞赛取得冠军的方案。考虑到OPPO手机的优异画质,还是觉得非常有必要研究一下这篇论文。

Abstract

由于不同图像的纹理细节变化极大,单帧图像极限超分难度很大。为解决该难度,作者提出了一种带感受野模块的ESRGAN超分网络,称之为RFB-ESRGAN。该文的主要贡献包含以下几个方面:

  • 为提取多尺度信息并增强特征判别性,作者采用感受野模块进行超分。感受野模块在目标检测和分类问题中取得了极具竞争力的结果。
  • 不同于多尺度感受野模块中的大卷积核,作者在感受野模块中采用多个小尺度卷积,这有助于提取细节特征并降低计算复杂度。
  • 在上采样阶段交替使用不同的上采样方法以降低计算复杂度,同时具有令人满意的性能。
  • 最后作者采用10个不同迭代的模型集成已提升模型的鲁棒性并降低单个模型引入的噪声。
  • 所提方法在NTIRE2020perceptual Extreme Super-Resolution竞赛中获得冠军。

Method

极限超分的目标在于:恢复高频细节同时 保持内容一致性。深度学习时代的超分方法往往采用PSNR作为训练损失,然后通过这样方式得到结果存在模糊现象,缺乏高频细节。感知驱动的方法有助于提升图像的视觉效果,比如SRGAN,ESRGAN等。

Basic Network Architecture

上图给出了所提方案的框架图,它包含五个模块:特征提取、Trunk-A、Trunk-RFB、上采样模块以及重建模块。

  • 特征提取:该模块比较简单,采用单个$3 \times 3$卷积构成;
  • Trunk-A模块由16个RRDB构成,RRDB如下所示,等同ESRGAN中RRDB;

image-20200527115004085

  • Trunk-B模块,如下所示。它与RRDB类似,区别在于它采用RFB提低RDB中的卷积。RFB由变尺度卷积构成,因此Trunk-B可以恢复图像的丰富细节。

image-20200527115130011

  • 上采样模块:它由RFB模块以及最近邻、pixelshuffle交替构成的上采样构成。
  • 复原模块:它由两个卷积构成。上采样与复原模块的整体结果如下:

image-20200527131643502

Multi-scale Receptive Field Block

为更好的恢复图像的细节,作者认为生成网络应该具有变尺度卷积(比如$1\times1, 3\times 3, 5\times 5$)。但是大尺度卷积会导致模型的推理时间变长,因此提出采用小尺度卷积替代。RFB最早被用于目标检测领域提升轻量型网络的特征表达能力,并取得了不错的性能。RFB充分利用多分支池化与卷积调整不同分支的感受野。

作者在原始RFB的基础上进行调整,最终采用的RFB如下所示。它主要有小尺度卷积构成,比如$1\times1,1\times3,3\times1,3\times3$。这种替代方式可以有效的降低模型的参数量与计算复杂度。与此同时,作者还提到:RFB的这种配置有助于提取细节特征(尤其是线条特征),比如头发、皮肤纹理以及边缘等,而这种功能也是超分网络最需要的功能。

Upsample Module

不同于现有超分网络采用单一上采样方式,比如双线性、转置卷积以及pixelshuffle。作者提出交替使用最近邻与pixelshuffle方式,最近邻执行空域变换,pixelshuffle执行depth-to-space变换。这种交替执行方式有助于提升space与depth间的信息交互,同时pixelshuffle有助于降低参数量和时间复杂度。

Loss Function

损失函数方面基本与ESRGAN相同,对于生成网络而言,损失函数由三部分构成:VGG内容感知损失、对抗损失以及像素损失($L1$损失)。对于判别网络而言,损失函数由两部分构成:$L{real}$促使真实图像与生成图像尽可能相似、$L_{fake}$促使生成图像与生成图像尽可能不相似。

Model Ensemble

在模型集成方面,ESRGAN对于PSNR驱动的模型与GAN驱动的模型在参数层面进行了融合。为更好的改进生成图像的视觉感知性能,作者在模型集成方面并未融合PSNR驱动模型;而是对训练过程中具有最好感知性能的不同迭代次数的10个模型进行融合(有点类似于EMA、Snapshot Ensembles)。融合方式与ESRGAN中的融合方式相同,故不再赘述。最终得到的集成模型可以有效减少重建图像的噪声,同时具有更好的鲁棒性。

Experiments

在训练阶段,训练数据采用matlab中的imresize获取,batchsize=16,输入LR图像大小为$32\times 32$。整个训练过程分为两个阶段,第一个阶段采用$L_1$损失进行训练;第二个阶段引入内容损失与对抗损失进行模型微调。注:这种分阶段的训练方式有助于避免GAN训练过程的不稳定问题。

训练数据集作者采用的DIV8K,它包含1500HR图像,1中1400用于训练,100用于验证;与此同时,作者还添加了DIV2K、Flickr2K以及OST数据。在数据增广方面,作者选用了随机水平镜像、随机90度旋转。

下图给出了RCAN、ESRGAN在DIV8K数据上的视觉效果对比。注:原始的RCAN与ESRGAN并不支持x16超分,作者对其进行微调并在上述数据集进行训练。从图中可以明显看出:所提方法生成的图像具有更锐利而清晰的纹理;而PSNR驱动的方法则倾向于生成平滑而模糊的图像,这与人眼视觉感知系统相悖。

与此同时,作者还对比了几种方法在DIV8K(1401-1500,未用于训练)数据上的指标对比,见下表。

image-20200527135718027

最后,作者还提供了竞赛结果对比,见下表。

image-20200527135835927

Ablation Study

为更好的说明所提网络不同单元的作用,作者进行了消融实验,见下图。第二列表示只用pixelshuffle上采样,第三列表示只用最近邻进行上采样,第四列表示采用两种交替方式进行上采样。

为更好说明RFB的效果,作者移除了所有RFB并保持网络结构不变,效果见第四列。可以看到:头发纹理细节明显变差且存在纹理方向错误。而第五列中使用RFB的模型则具有更好的头发纹理细节。

为说明本文所提交替上采样方案的有效性,见上图第2-4列。可以看到:pixelshuffle上采样结果明显优化最近邻上采样;而仅仅使用pixelshuffle生成的图像则过于锐利并不自然,同时还会生成假性伪影;而交替上采样方式则具有更好清晰而真实的结果。

为说明模型集成的有效性,见上图第5-6列。模型集成有助于降低噪点同时生成的纹理更为自然。

最后,作者还给出了上述不同配置模型的DIV8K(1401-1500)数据上的指标对比,见下图。

image-20200527141244400

Conclusion

本文对NTIRE2020感知极限超分竞赛冠军方案进行了解决,它在ESRGAN的基础上进行了几点改进:

  • 采用多尺度RFB提取LR图像的多尺度特征信息;
  • 提出了最近邻与pixelshuffle交替上采样的方案;
  • 为抑制生成图像的噪声,提出了不同阶段模型融合方案。

Appendix

了解ESRGAN的小伙伴复现起该论文应该不是什么难题,感兴趣的可以尝试一下哈(注:如有兴趣也可以与Happy一起讨论代码实现与复现哦)。另外,比较有意思的一点:纵观NTIRE2020竞赛方法,可以看到不少ESRGAN的身影。

  • 在感知极限超分竞赛中,冠军方案OPPO-Research基于ESRGAN改进而来;
  • 在真实图像超分竞赛中,冠军方案也是基于ESRGAN改进而来(网络架构完全相同、数据不同、判别部分有点区别)。除了两个冠军方案,还有其他参赛方案也是在ESRGAN的基础上进行的改进。

感兴趣的小伙伴可以去看一下这两篇超分竞赛方法总结:

NTIRE2020 Challege on Real-World Image Super Resolution: Methods and Results. https://arxiv.org/abs.2005.01996.

NTIRE2020 Challege on Perceptual Extreme Super Resolution: Methods and Results. https://arxiv.org/abs/2005.01056.

微信公众号: 极市平台(ID: extrememart )
每天推送最新CV干货