基于递归金字塔的多任务轻量化图像超分辨率算法

(重庆理工大学计算机科学与工程学院,重庆 400054)

1 本文方法

1.1 整体网络结构

我们提出的递归金字塔多任务网络(Recursive Pyramid Multitasking Networks,RPMN)可以在单个网络中端到端地解决多个倍数的放大任务,是一个的轻量化多任务网络。如图1所示,本文算法RPMN 由三个部分组成,分别是浅层特征提取模块,递归金字塔,重构模块。

图1 算法流程Fig.1 Algorithm flow

1.2 递归金字塔模块

为解决超分辨率网络层数需求大和网络单任务的问题,我们构建一个轻量化递归多任务网络,其中最重要的模块便是我们设计的递归金子塔模块。

我们整体网络结构受递归网络[1]的启发,设计了递归金字塔模块。整体的设计思想是:将一个2n的放大任务分解为n 个2 倍放大任务,将2 倍放大任务以相同的参数建模,使用递归结构的形式在每一次的递归中完成,n 次递归便能完成2n倍数的超分任务。

基于上述的设计思想,我们设计了轻量化多任务的递归金字塔模块。如图2所示。

图2 递归金字塔模块Fig.2 Recursive Pyramid Module

图2 是我们算法的2 层递归金字塔。该网络结构受到DRCN和多任务学习的启发,绿色框中的部分采用作为迭代的结构整体,所有卷积和反卷积(粉红部分)共享卷积,与DRCN相比做了以下3点改进:(1)独立的重建卷积解决单网络多任务的问题;(2)在端到端网络里同时训练多个任务,各任务间正向促进;(3)使用了LAFM替换了普通的卷积堆叠,充分地轻量化网络,提高特征提取能力。网络的各级特征图呈现出金字塔的形状,故将网络称为递归金字塔。

1.3 局部自适应融合模块

算法中非映射主体是递归金字塔,递归金字塔主要由多个局部自适应融合模块构成(Local adaptive fusion module,LAFM),LAFM设计如下图3所示。

我们首先定义xn-1和xn为第n个LAFM的输入和输出特征图。LAFM的公式化表示为式(1):

为解决特征提取效率问题,本文设计的LAFM模块便是针对此问题设计的轻量化高效特征提取模块。LFAM由多个密集连接注意力模块(Densely connect attention modules,DCAM)构成,其轻量程度取决于DCAM。DCAM主要由密集连接[2]与通道注意力机制[3]构成,如图4所示。

将每4 个卷积看为一个整体,4 个卷积间采用密集连接,每一个卷积的输入会将之前所有卷积的输出特征复用,在训练阶段的反向传播使第N层的权重不仅与第N-1层相关,还能调整第N 层以后的所有权值。通道注意力则是更加关注于特征中的高频信息,提高特征提取质量。

LAFM 在DCAM 的基础上考虑了DCAM 模块间的特征融合,从而提高特征表征能力。主要有以下两个方面的改进:(1)添加少量分支权重分配参数λ,用于自适应分配各DCAM输出特征的融合权重。(2)采用1×1卷积可以融合各DCAM 的信息。

2 实验与分析

2.1 数据与评价指标

数据集:DIV2K是2017年Ntire图像超分辨率挑战赛的公开数据集,共计1000张2k分辨率的高质量图像,我们将DIV2K的前800训练集。我们使用的测试集是公开的4个数据集:Set5、Set14、B100、Urban100。

图3 局部自适应融合模块Fig.3 Adaptive fusion module

图4 DCAM 模块细节Fig.4 DCAM module details

图5 ×4 超分任务的可视化对比Fig.5 ×4 Visual comparison of super-resolution tasks

验证标准:我们将超分率图像转换到YcbCr 颜色空间,测试亮度(Y)通道上的峰值信噪比PSNR和结构相似度SSIM。

2.2 实验结果

本文算法与Bicubic、SRCNN、FSRCNN、VDSR、DRCN、LapSRN、MSRN、AWSRN 进行了定性和定量的对比。在Set5、Set14、B100、Urban100上分别对比了他们的PSNR和SSIM。最后在×4的超分任务中计算了参数量和运算量,便于进行轻量性的对比。定量对比表格如表1所示。

表1 ×4 超分任务的不同算法定量对比Tab.1 ×4 Quantitative comparison of different algorithms for super-resolution tasks

除本文算法外,其余算法的×4模型都是分开训练的独立模型。观察表1可以发现,本文算法在4个公开数据集上的PSNR 和SSIM 中均取得最优重建效果。在参数量和运算量方面,参数量约减少至AWSRN 的三分之一,四个数据集上的PSNR也分别得到了0.13db、0.24db、0.17db、0.41db的提升。

从定性的角度分析,观察可视化图5,钢架的重建图,我们可以清楚地发现本文算法重建的钢条结构更加笔直和清晰,与真实图片最接近,明显优于其他算法。

3 结论

针对主流超分算法网络只能解决单任务的问题,本文对递归网络进行改进,对每次递归的输出进行独立的特征图重建,实现一个多任务网络。同时在局部特征提取模块的设计上采用密集连接和特征自适应融合,既保证特征提取能力又达到了网络轻量化。与主流的CNN 方法的图像超分辨率算法进行了对比,在模型参数,模型计算量,重建精度上均取得了不错的提升。该算法对于现实中的超分辨率应用具有重要意义。