RepRFN:当RFDN遇到重参数化
0. 简介
NTIRE 的全称为New Trends in Image Restoration and Enhancement Challenges
,即“图像复原和增强挑战中的新趋势”,是CVPR(IEEE Conference on Computer Vision and Pattern Recognition)举办的极具影响力的计算机视觉底层任务比赛,主要涉及的研究方向有:图像超分辨率、图像去噪、去模糊、去摩尔纹、重建和去雾等。
其中在2023年,CVPR开展的NTIRE相关挑战有:
-
夜间摄影渲染(night photography rendering);
-
从图像中估计深度(HR depth from images of specular and transparent surfaces);
-
图像去噪(image denoising);
-
视频着色(video colorization);
-
阴影去除(shadow removal);
-
视频增强的质量评估(quality assessment of video enhancement);
-
双目超分辨率(stereo super-resolution);
-
光场图像超分辨率(light field image super-resolution);
-
图像4倍超分辨率(image super-resolution (×4));
-
360°全景图像和视频超分辨率(360° omnidirectional image and video super-resolution);
-
镜头到镜头的散景效果转换(lens-to-lens bokeh effect transformation);
-
实时4K超分辨率(real-time 4K super-resolution);
-
高分辨率图像非均匀去雾(HR nonhomogenous dehazing);
-
高效超分辨率(efficient super-resolution)。
同时,以上的这些挑战也蕴含着当前的一些研究难点及挑战,需要研究学者们集思广益,提出针对提升任务性能的想法,为共同解决近年来的难题贡献出一份力量。
本篇文章着重于NTIRE 2023高效超分辨率(efficient super-resolution)挑战赛的冠军方案的解读,学习能够提升任务的tricks,以期给相关的科研任务一些启发。NTIRE 2023高效超分辨率挑战赛的目标是以RFDN(AIM 2020 高效超分辨率冠军)为baseline,以尽量少的推理时间(runtime),参数量(parameters),计算量(FLOPs),激活值(activations)和显存占用(memory consumption)实现4倍超分,同时在DIV2K validation set上的PSNR至少要达到29.00dB。
比赛提供的数据集有DIV2K数据集和LSDIR数据集。DIV2K数据集包含1000张多样的2K分辨率的RGB图像,其中训练集有800张,验证集有100张,测试集有100张。LSDIR数据集有86991张高分辨率高质量的图像,其中训练集有84991张,验证集有1000张,测试集有1000张。
1. 摘要
为了解决在资源受限的设备上很难部署超分模型(模型参数、计算量过大)的问题,本文探索了轻量化超分中信息蒸馏机制(information distillation mechanism)和残差学习机制(residual learning mechanism) 在性能和效率上的差异,提出了基于重参数化(reparameterization) 的轻量级超分网络结构,叫做RepRFN。RepRFN能够高效地减少GPU占用和提高推理速度。
本文提出了一个多尺度特征融合结构(multi-scale feature fusion structrue) ,能够使得网络学习和聚合不同尺度的特征和具有高频信息的边缘。作者重新思考了整个网络框架存在的冗余性,尽可能移除一些不影响性能的冗余模块,从而降低模型的复杂度。此外,作者引入了一个基于傅里叶变换的损失函数,将图像从空间域变换到频率域,使得网络能够监督和学习图像的高频信息。
论文代码链接:https://github.com/laonafahaodange/RepRFN
2. 引言
近年来,许多基于CNN的SR网络被提出,也意味着CNN对图像SR的发展起着重要的作用。
- 2014年,Dong等人提出了第一个基于CNN的超分方法——SRCNN。
- Kim等人提出了一个更深(20层)的网络来提高超分性能——VDSR。
- Lim等人使用局部和残差连接,提出了EDSR。
然而,大多数SR网络为了提高图像复原时的性能往往会牺牲效率。在某些情况下,实时性也会影响用户的体验。因此,如何高效地提取图像的边缘,纹理,结构和其他信息,同时平衡SR网络的性能和复杂程度是至关重要的研究,这决定了网络是否可以部署在资源受限的设备上。
针对上述情况,论文提出了一个重参数化残差特征网络(Reparameterized Residual Feature Network) ,也叫做RepRFN。作者设计了一个多分支结构,通过使用多个并行的不同尺寸的卷积核,提取不同感受野的特征,并使用局部残差连接(local residual connection) 进行特征融合(feature fusion)。为了高效提取边缘信息,在多分支结构中使用了边缘导向卷积块(Edge-oriented Convolution Block, ECB)中的Sobel分支和Laplace分支。在训练阶段,作者将SR任务看作是空间域学习和频域学习的多任务学习问题,使用基于傅里叶变换的损失函数,指导模型复原高频信息。实验表明,提出的RepRFN在性能和效率上取得了很好的均衡。
作者的贡献总结如下:
- 首先,提出了基于重参数化的多尺度特征融合结构,通过多个并行的不同感受野的卷积和边缘导向卷积模块提取不同模型的特征,使用残差连接聚合这些特征,提高模型的特征表达能力;
- 重新考虑了RFDN模型的结构,分析RFDN的冗余性,在作者的网络中移除了用于通道变换的 1 × 1 1\times1 1×1的卷积;
- 引入了基于傅里叶变换的损失函数,使得模型能够在训练过程中学到图像的频率信息,增强模型恢复频率细节的能力。
3. 相关工作
相关工作中回顾了一些主流的高效图像超分领域的方法。
- Dong等人提出了第一个基于CNN的超分方法——SRCNN。
- 比SRCNN速度快17倍——FSRCNN。
- Kim等人提出了一个深度递归卷积网络——DRCN。
- 在DRCN的基础上,结合递归和残差网络,DRRN诞生。
- 2018年,NamhyukAhn等人使用群卷积(group convolution)提高网络效率,在级联模块之间使用类似于递归网络的机制共享参数,提出了轻量级级联残差网络——CARN。
- Lai等人结合传统图像算法拉普拉斯金字塔与深度学习,实现多级超分辨模型——LapSRN。
- Hui等人提出了一个信息蒸馏网络——IDN。
- 基于IDN,信息多蒸馏网络IMDN由一系列级联的IMDB块组成,IMDN获得了AIM2019资源受限超分挑战赛的冠军。
- Liu等人重新思考IMDN,提出了残差特征蒸馏网络——RFDN。
- E-RFDN获得了AIM2020高效超分挑战赛的冠军。
4. 方法
-
在4.1节中,作者提出了残差特征网络,与信息蒸馏机制对比,作者通过实验观察到残差特征机制和信息蒸馏机制的差异。
-
在4.2节中,作者回顾了残差特征网络的缺点,提出了一个多尺度特征融合的基于重参数化的轻量级SR网络RepRFN。
-
在4.3节中,作者介绍了基于傅里叶变换的损失函数,将图像从空间域变到频率域,使得模型能够在训练过程学习到频率信息。
4.1 残差特征网络(RFN)
残差特征网络的结构如下图,包含浅层特征提取模块,深层特征提取模块和上采样模块。(这样的结构实际上已经是超分的经典结构了,比如SwinIR、HAT都基本由这三个模块构成,只不过各个模块内部的细节可能不一样,大部分人改进的地方都在深层特征提取模块中)浅层特征提取模块用来提取LR图像的浅层特征,深层特征提取模块对浅层特征进行进一步的非线性映射,得到深层特征。然后,通过残差连接将浅层特征和深层特征进行融合。最后,上采样模块对融合后的特征进行重新组合,得到重构的SR图像。
看上面的图,我们很容易知道:
- 浅层特征提取模块由一个 3 × 3 3\times3 3×3的卷积层组成。
- 深层特征提取模块由一组堆叠的残差特征块组成,能够逐步提取浅层特征,使用残差连接去集成浅层特征和深层特征,提高模型的特征表达能力。
- 上采样模块由一个 3 × 3 3\times3 3×3的卷积层和一个PixelShuffle层组成。
残差特征块的关键在于残差特征学习机制。信息蒸馏机制是将输入特征沿着通道维度分为两个部分,一个部分保持不变,另外一个部分输入到下一个信息蒸馏模块中,进一步提取特征。经过几个蒸馏步骤之后,会沿着通道维度进行concatenate,从而完成特征融合操作,这样就实现了蒸馏信息的融合。然而,残差特征学习机制不一样,它不是沿着通道维度划分特征,而是直接将提取的特征输入到下一个模块,只是简单的add和merge每个模块提取的深层特征和浅层特征,这样能够避免占用显存过大的问题(这些问题往往是通道划分和concatenate操作造成的),加速推理时间。下图展示了几种不同的信息蒸馏模块,可以看到本文用到的残差特征块(RFB) 不是像RFDN-IDB那样使用通道划分,而是直接输入到下一个卷积层,使用残差融合代替了信息融合机制。
同时,作者探索了信息蒸馏机制和残差学习机制在性能和效率之间的差异。下图的RFB1代表的是局部残差连接,RFB2代表的是全局残差连接,RFB3代表的是局部和全局残差连接的组合。Attention Layer使用的是和RFDN相同的增强空间注意力(ESA)。将没有任何残差连接的模型作为baseline模型。
从下表可以看出,全局残差连接带来的增益会少于局部残差连接。
4.2 重参数化残差特征网络(RepRFN)
3 × 3 3\times3 3×3的卷积层通常用来提取特征,但是它的感受野太小了。其次,RFN模型的结构还是存在冗余。另外,在图像特征域提取和恢复高频信息仍有不足。因此,作者对RFN模型再进行改进,提出了一个多尺度特征融合的基于重参数化的轻量级SR网络RepRFN。
为了解决感受野太小的问题,作者设计了多并行分支结构,提取和融合不同感受野和模式的特征,使得模型尽可能受益于多分支结构。同时,重参数化操作将训练和推理过程进行了解耦,避免了由于引入多分支结构导致的参数量和计算量增大的问题。
为了解决模型结构冗余的问题,作者重新思考和分析了RFN和RFDN之间的结构差异,移除了RFN中用于通道变换的 1 × 1 1\times1 1×1的卷积层,对ESA进行了结构上的改进。
RepRFN和RFN有相同的结构,区别在于使用RepRFB替代了图2中的RFN。RepBlock是RepRFB的主要组成部分,而多分支结构组成了RepBlock,如下图所示。
RepRFB的设计参考的是RFDN中的RFDB,在RFDB中,中间特征图在每个信息蒸馏模块中会被浅层残差块(SRB)划分三次(如图3c所示),因此RepRFB中的首个三层采用的就是一个重参数化的多分支结构,在本文中也称作RepBlock,特征通过执行不同操作的路径进行传递,最后将它们融合在一起,以期提高模型的表达能力。在RepRFB中,由于存在局部残差连接,在RepBlock和卷积层前后的中间特征图尺寸和通道数量通常是不变的,所以不需要通道变换操作,因此RepRFB中的 1 × 1 1\times1 1×1卷积是冗余的,移除该卷积能进一步减少参数量。
4.3 基于傅里叶变换的损失函数
针对提取和复原高频信息的问题,除了引入多分支结构的ECB以外,还引入傅里叶变换到损失函数中,指导模型尽可能学习频域特征和复原高频信息。基于傅里叶变换的损失函数如下:
L f = ∣ ∣ f f t ( I S R ) − f f t ( I H R ) ∣ ∣ 1 L_{f}=||fft(I_{SR})-fft(I_{HR})||_1 Lf=∣∣fft(ISR)−fft(IHR)∣∣1
对应该损失函数的核心代码如下:
fft_loss = self.l1loss(torch.fft.fft2(X, dim=(-2, -1)),torch.fft.fft2(Y, dim=(-2, -1)))
5. 实验
5.1 实验设置
训练集:DIV2K和Flickr2K。HR图块大小为 192 × 192 192\times192 192×192。
数据增强策略:随机水平、垂直翻转和旋转。
优化器:Adam。 β 1 = 0.9 , β 2 = 0.999 \beta_1=0.9, \beta_2=0.999 β1=0.9,β2=0.999。
测试集:Set5、Set14、BSD100、Urban100和Manga109。
训练策略:初始学习率设置为 5 × 1 0 − 4 5\times10^{-4} 5×10−4,每100个epoch学习率减半,一共训练1001个epoch。
5.2 客观结果
PSNR和SSIM是在Y通道上进行计算的。计算参数量和运算量时,假设模型输出的图片是720P。
可视化结果比较。
5.3 消融实验
多尺度特征融合模块
其中RepRFN-P代表没有使用多分支结构(P代表Plain)
模型结构
为了获得低复杂度的模型,作者牺牲了一些性能来换取低复杂度,最后的RepRFN模型使用48个通道数,改进的ESA模块,并移除了用于通道变换的 1 × 1 1\times1 1×1卷积。
损失函数
论文细节问题,表格标题最后面的RGB color soace
应该是RGB color space
。
5.4 NTIRE 2023高效超分挑战赛
6. 结论
在本文中,作者提出了一个重参数化残差特征网络用于轻量级图像超分,设计了多分支结构用来尽可能捕获不同模式的特征,并将这些特征融合。其次,引入了重参数化操作使得复杂的多分支结构也能够用于轻量化网络中。在网络训练过程中,设计了基于傅里叶变换的损失函数,能够将空间域变换到频率域来指导模型学习频率信息。实验表明所提的方法相比于其他网络能够在性能和效率间达到更好的均衡。
最后感谢小伙伴们的学习噢~
最后附上2023年高效超分的比赛报告链接,欢迎大家多多阅读分享:NTIRE 2023 Challenge on Efficient Super-Resolution: Methods and Results