基于自校准双注意力的图像超分重建网络

周华平,郭颍杰,孙克雷

(安徽理工大学计算机科学与工程学院,安徽 淮南 232001)

给定一张低分辨率(Low Resolution,LR)图片重建出对应的高分辨率(High Resolution,HR)图片的过程,称为单图像超分辨率(Single Image Super Resolution,SISR)。重建后的图像视觉质量较高,纹理细节丰富,这为克服医学成像[1-2]、卫星成像[3]、安全与监控等各种计算机视觉应用中的分辨率限制提供了条件。得益于卷积神经网络强大的特征表达能力,构建一个端到端的卷积神经网络模型,学习LR图片与HR图片之间的非线性映射函数,成为目前主流的研究方法。

文献[4]提出开创性工作SRCNN(super-resolution using deep convolutional networks),使用插值技术对LR图像进行预处理,上采样至目标尺寸作为模型的输入,经过三层卷积提取图像的特征。文献[5]构建了一个极深的网络结构VDSR(very deep convolutional networks),通过构建更深的网络层获得更大的感受野,并利用自适应梯度裁剪技术加快收敛过程。文献[6]提出DRCN(deeply-recursive convolutional network),使用参数相同的递归块来代替不同的卷积层,增加网络深度同时降低参数量。文献[7]提出深度持续记忆网络(Deep Persistent Memory Network,Mem Net),利用记忆模块多次递归,更好地利用多层特征信息。然而,上述模型通过对输入特征进行放大处理,这会增加模型计算复杂度并产生可见的伪影。

文献[8]提出的EDSR(enhanced deep residual networks)获得了NTIRE2017 SR挑战赛的冠军,通过移除批归一化层以及使用更深更宽的网络结构来增加网络表达能力。文献[9]结合密集连接与残差连接提出了RDN(residual dense network),充分利用层次结构特征。文献[10]提出RCAN(residual channel attention networks),通过提高对通道间重要信息的关注度,使得网络可以恢复更多的高频信息,从而提升模型的重建效果。然而,上述算法为了获得更好的重建性能,不断加深网络的层数,不可避免地导致参数量与计算量大幅度增加。此外,随着网络层数逐渐加深,高层次的特征更倾向于表达图像中低频的语义信息,导致对提高超分性能至关重要的纹理细节缺失。为了解决这类问题,提出了一种基于自校准双注意力的图像超分重建网络(image super-resolution reconstruction network based on self-calibrated and dual-attention,SCDAN),主要贡献如下:

(1)设计了一种自校准残差块(self-calibrated residual block,SCRB),使用校准分支和残差块分支分别提取特征信息。其中,校准分支使用像素注意力(pixel attention,PA)为特征图中的所有像素点生成注意力系数,有目的地恢复图像纹理信息。残差块分支利用跳跃连接将低频信息与输出特征相加,减缓梯度消失,稳定训练。最后,对两个分支的特征进行融合以提升模型的表达能力。

(2)在由多个SCRB构成的自校准残差组(self-calibrated residual group,SCRG)的组内与组间进行分层特征融合,充分利用模型的浅层与深层信息,缓解由于网络逐渐加深导致特征在传播过程中信息缺失的问题。

(3)提出了一种双注意力机制(dual-attention,DualA),结合通道注意力与空间注意力增强模型在通道和空间维度的建模能力,对融合后的信息进行过滤,保留图像中的高频细节。

1 自校准双注意力图像超分重建网络

1.1 网络整体结构概述

自校准双注意力网络(Self-calibrating Dual Attention Networks,SCDAN)主要通过级联非线性映射模块和图像重建模块实现高质量超分重建,其结构如图1所示。定义网络的输入图像为ILR,输出图像为ISR。

图1 SCDAN整体网络结构

网络首先使用一个3×3大小的卷积层将输入图像ILR∈H×W×3(H、W分别表示图像的高和宽)映射到高维空间,其操作可表示为

F0=f3×3(ILR)

(1)

式中,f3×3表示3×3卷积操作;F0表示浅层特征提取模块的输出特征。

接着将F0送入非线性映射模块,进一步拟合输入图像ILR与输出图像ISR之间的非线性映射关系。

(2)

式中,fDualA表示双注意力机制;f1×1表示1×1卷积操作;[·]表示特征通道级联操作;Gi-1表示第i个SCRG输入。则非线性映射模块可表示为

FN=fDualA(f1×1([G1,G2,…,Gm]))

(3)

式中,FN表示非线性映射模块的输出特征。通过对SCRG的组内与组间引入分层特征融合结构,可以缓解随着网络加深,模型逐渐丢失部分纹理细节信息的问题,增强模型对特征的复用能力。融合后的特征可能包含有大量的冗余信息,通过设计的双注意力机制对信息进行过滤,使得网络更加关注特征通道与空间中的重要信息,可以恢复对图像超分至关重要的高频信息,获得令人满意的重建结果。

图像重建模块将浅层特征提取模块的输出F0与非线性映射模块的输出FN相加作为输入特征,通过亚像素卷积层进行上采样,再经过一个3×3卷积层调整通道维度得到输出图像ISR。其操作可表示为

ISR=f3×3(fPixelShuffle(F0+FN))

(4)

式中,fPixelShuffle表示亚像素卷积操作。

1.2 自校准残差块

LR图像中含有大量易于学习的低频信息,利用残差学习可以将这些信息通过恒等映射直接传输到网络末端,使网络直接学习对图像超分任务至关重要的高频信息,缓解梯度消失问题。文献[11]提出一种自校准卷积,将输入特征进行通道拆分,以并行的双分支结构对特征进行特征提取和重校准。结合残差块与自校准卷积的优点,提出了一种自校准残差块(SCRB),利用传统的残差块作为其中一个分支,用于捕获LR图像的特征信息。另一分支使用像素注意力(PA)为特征图所有像素点生成注意力系数,有目的地恢复图像的纹理信息。

如图2所示,设计的自校准残差块由并行的残差块分支和校准分支构成。输入特征Xn-1在经过两条支路前,会经过两个1×1卷积层进行降维,将通道数减半以减少模型参数量,同时得到两条支路的输入H1和H2。其操作可表示为:

图2 自校准残差块

(5)

(6)

H1,H2首先分别经过残差块分支和校准分支,然后对两个分支的输出特征进行特征融合并与输入特征Xn-1相加,得到输出特征Xn。其操作可表示为:

(7)

F2=f3×3(H2)⊙σ(f1×1(H2))

(8)

Xn=f1×1([F1+F2])+Xn-1

(9)

式中,Φ表示LReLU激活函数;σ表示Sigmoid激活函数;⊙表示逐元素相乘操作;F1,F2分别为残差块分支和校准分支的输出。

1.3 双注意力机制

通过分层特征融合操作,可以充分复用各层级的输出特征,但是会带来大量冗余信息。为了提高网络提取图像高频细节的能力并减少冗余信息对后面网络层的影响,SCDAN在组内与组间的分层特征融合操作后加入一个双注意力机制(DualA),对融合后的特征进行信息过滤,使得网络更加关注通道与空间的重要特征,忽略那些不重要的特征。

DualA结构如图3(a)所示,主要采用通道注意力(channel attention,CA)和空间注意力(spatial attention,SA)并行的模式分别处理通道维度和空间维度上的重要信息。Transformer中广泛使用的多头自注意力机制可以建立像素间的长程依赖关系,但是其计算量随输入大小呈二次增长。受到文献[12]启发,通过在特征的通道维度而非空间维度上计算像素间的相关性,可以降低模型的计算负载,计算量可从O(W2H2)降至O(C2)。对于通道注意力分支(CA),其结构如图3(b)所示。给定输入特征Xn-1,CA分支首先使用3个1×1卷积生成查询向量(Q)、键向量(K)和值向量(V)。其操作可描述为:

图3 双注意力机制组件

(10)

(11)

(12)

通过reshape操作调整向量的维度得到新的Q′∈C×HW、K′∈C×HW和V′∈C×HW。然后,将Q′与K′的转置执行点乘操作并经过Softmax函数生成权值图A∈C×C。最后,将权值图与V′相乘,得到加权后的特征FCA,其操作可表示为

FCA=Softmax(Q′·(K′)T/α)·V′

(13)

式中,Softmax表示用于生成权值图的Softmax函数;α是一个可学习的缩放参数,用于控制Q′和K′的点积的大小。通过对特征通道间的长程依赖关系进行建模,可以使网络能够捕获远距离相似纹理,提升模型的特征表达能力。

为了进一步恢复图像精确的空间结构,空间注意力分支(SA)首先对输入特征在通道维度上使用平均池化和最大池化生成两个大小为1×H×W的特征图,然后对两个特征图按通道进行拼接,并通过卷积核大小为7×7的卷积层提取特征图的空间信息,最后经过Sigmoid函数生成权值图并对输入特征进行加权,其结构如图3(c)所示。该过程可表示为

FSA=Sigmoid(f7×7([AvgPool(Xn-1);MaxPool(Xn-1)]))⊙Xn-1

(14)

式中,AvgPool(·)和MaxPool(·)分别表示平均池化和最大池化操作;f7×7表示7×7卷积;FSA表示SA分支的输出。将两个分支的输出特征图进行融合并通过1×1卷积进一步调整,得到输出特征Xn,使得模型能够更好地恢复图像的高频信息。其操作可表示为

Xn=f1×1(FCA(Xn-1)+FSA(Xn-1))

(15)

通过结合通道注意力分支与空间注意力分支,增强了模型在通道和空间维度的建模能力,提升了模型的判别能力,使得模型重建结果更加清晰,纹理细节更加丰富。

1.4 损失函数

(16)

式中,fSCDAN为本文提出的网络模型;θ为网络中的可训练参数。

2 实验

2.1 数据集与评价指标

本文实验在公共数据集DIV2K[13]数据集和Flickr2K[14]上进行训练,共包含3 450张对应的低分辨率图片(LR)和高分辨率图片(HR)。其中,在DIV2K数据集中选取前800张和Flickr2K组成实验训练集。实验测试集包含:Set5[15]、Set14[16]、BSD100[17]、Urban100[18]。本文将峰值信噪比(PSNR)和结构相似度(SSIM)作为性能评价指标,并在YCbCr通道的Y通道(亮度)进行计算。

2.2 参数设置

在模型训练时采用Adam算法对模型进行优化。批大小设置为16。数据增强使用随机旋转90°、180°、270°和水平翻转。本文对2×,3×和4×三种缩放因子进行实验,将大小为48×48的LR子图作为模型的输入,对应的真实标签HR在3种缩放因子下的大小为96×96、144×144和192×192。实验初始时,学习率为5×10-4,每2×105个迭代学习率减半,总迭代次数为8×105。模型使用的深度学习框架为Pytorch,编程语言为Python3.8,在NVIDIA RTX 3090进行训练和测试。

2.3 实验结果与分析

为了验证本文提出结构的有效性,进行了一系实验。为了公平比较,所有实验在同一环境下,迭代次数为5×105。

1)自校准残差块有效性分析 自校准残差块(SCRB)是非线性映射模块的核心结构,为了验证SCRB的有效性,本文与传统的残差块(ResBlock)结构进行了对比实验,结果如表1所示。其中,“ResBlock”表示用传统的残差块替代SCRB。相比于传统的残差块,使用SCRB将PSNR平均提高了0.03dB,而参数量与计算量大大降低。这得益于校准分支使用像素注意力(PA)为特征图所有像素点生成了权重信息,可以有目的地恢复图像纹理信息。

表1 自校准残差块(SCRB)性能分析

2)特征融合与双注意力机制有效性分析 为了证明设计的双注意力机制与特征融合结构的有效性,进行了5组消融实验,结果如表2所示。其中,“FF”表示特征融合,“CA”为通道注意力分支,“SA”为空间注意力分支。由表2中前2组消融实验对比分析可得,在自校准残差组的组内与组间进行分层特征融合在4个测试集上的PSNR平均提升了0.047 5dB。这表明通过聚合浅层与深层特征,可以缓解因网络不断加深导致的信息丢失,充分利用各层特征,提升模型的表达能力。

表2 特征融合与双注意力机制性能分析

表2中后3组实验主要对双注意力机制的有效性进行评估。通过表2结果分析,在使用分层特征融合的基础上,单独使用通道注意力分支和空间注意力分支分别将4个测试集上的PSNR平均提升了0.075dB和0.06dB。这表明使用注意力机制可以对聚合后的信息进行过滤,更加关注那些对于图像恢复至关重要的高频信息。双注意力机制同时关注特征的通道和空间维度的重要特征,在4个测试集上的PSNR性能指标平均提高了0.15dB,超越了单独使用通道注意力或空间注意力,充分体现了本文设计的双注意力机制在处理图像高频信息的有效性。

3)复杂度分析 本文网络在Set5数据集4×放大倍数下的PSNR与不同网络的参数量(Params)与浮点量(Flops)对比如表3所示,比较网络包括VDSR、DRCN、CARN[19]、LAPAR-A[20]、IMDN、LBNet[22]。对于轻量化网络来说,可训练参数的数量是需要考虑的基本元素。VDSR虽然参数量要小于本文算法,但是该算法的输入与输出大小相等,导致模型计算量急剧增加。相较于DRCN、CARN,本文模型的参数量仅为这两种模型的49%和55%,但PSNR却分别提高了0.9dB和0.3dB。IMDN与LBNet虽然参数量与计算量略小于本文算法,但是本文算法的重建性能更有优势。图4对比了本文网络与其他轻量级网络在Set5数据集2倍放大上的参数量与PSNR值的对应关系,本文所提网络保持较低参数量的同时,取得了最好的性能。这表明本文所提网络可以在性能和参数量上取得很好的权衡。

表3 4×放大倍数的复杂度分析

图4 Set5数据集2×放大倍数模型参数量和PSNR值

2.3 与其他轻量级网络的比较

1)定量客观比较 将本文所提网络与其他轻量级网络在2×、3×、4×放大倍数上进行了定量客观比较,包括:SRCNN、VDSR、DRRN[23]、DRCN、MemNet、IDN[24]、CARN、IMDN、LAPAR-A、RFDN[25]、LBNet、SCET[26]。实验对比结果如表4所示,所提网络在4个测试集上PSNR和SSIM性能指标均达到最优,尤其在Urban100数据集上性能提升较为明显,该数据集中主要含有大量的城市图片,包含大量的边缘信息且细节分布在不同频段,对于图像超分任务来说十分具有挑战性。本文网络在Urban100数据集2×、3×、4×放大倍数对比参数量最多的DRCN,PSNR分别提高1.64dB、1.31dB、1.21dB。这主要归因于提出的双注意力机制,其中,CA分支利用自注意力机制对全局相关性进行建模,建立图片中像素间的远距离相关性,而城市建筑多为重复性结构,所以可以更好地恢复远距离相似纹理特征。同时,SA可以捕获图像中的空间位置信息,可以更好地恢复复杂的结构特征。

表4 在四种数据集上,不同网络在2×、3×和4×上的平均PSNR和SSIM定量比较

2)主观效果展示与分析 本节详细分析了所提网络与其他轻量级网络的可视化结果,包括:Bicubic、SRCNN、IDN、LAPAR-A、CARN和IMDN。图5展示了在放大倍数为4时,不同网络对Set14数据集中的图片“Barbara”的重建效果。可以看出,其他网络恢复的书籍边缘均出现了扭曲、重叠,本文网络重建的书籍边缘更加锐利。图6展示了在放大倍数为4时,各种网络对BSD100数据集中“253027”的重建效果。其他网络重建的斑马条纹较为模糊,而本文所提网络重建效果较为清晰且恢复了更多的高频细节。

图5 放大倍数为4时,不同网络对Set14数据集中的“barbara”图像局部放大效果可视化对比

图7、图8展示了在放大倍数为4时,本文网络与其他算法对Urban100数据集中的两张图片“img_005”和“img_096”的重建效果。双三次插值、SRCNN生成结果比较模糊,失去了大部分结构信息。IDN、LAPAR-A、和IMDN对于墙体的边缘恢复较差,出现不同程度的模糊、扭曲,无法恢复锐利的纹理细节。相比之下,本文网络可以恢复更多的细节和更清晰的边缘,并获得更高的视觉质量,这归功于更有效的特征提取和特征复用。

图7 放大倍数为4时,不同网络对Urban100数据集中的“img_005”图像局部放大效果可视化对比

图8 放大倍数为4时,不同网络对Urban100数据集中的“img_096”图像局部放大效果可视化对比

4 结论

在保持模型性能的基础上减少模型的参数量与计算量,是SR领域热门研究方向。本文提出了一种基于自校准双注意力的图像超分重建网络,设计了一种自校准残差块,利用自校准的思想有目的地恢复图像纹理信息。将特征融合结构与设计的双注意力机制相结合,充分复用各层级的输出特征,同时过滤冗余特征,更加关注融合后的特征在通道和空间维度的重要信息,以恢复图像缺失的高频信息。在4个测试集上的定量与定性分析表明,与其他轻量级方法相比,模型在视觉质量和客观指标方面均取得不错地结果。同时,所提算法在模型性能与复杂度之间取得了很好的平衡。