基于双注意力机制的轻量级图像超分辨率重建

冯兴杰 王 荣

1(中国民航大学信息网络中心 天津 300300) 2(中国民航大学计算机科学与技术学院 天津 300300)

0 引 言

在实际应用中,受到图像采集设备成本、传输带宽、存储年限久远等问题的限制,很难保证得到的图像均为精致的高分辨率图像。因此超分辨率重建(Super-Resolution,SR)技术的产生显得尤为重要。根据输入数据类型的不同,可以将SR分为单图像SR和多图像SR,由于多图像SR直接基于单图像SR,因此本文主要讨论单图像超分辨率重建(Single Image Super-Resolution,SISR)。

图像超分辨率重建的目的是将给定的具有粗糙细节的低分辨率图像(LR)转换为相应的具有更好的视觉效果以及精致细节的高分辨率图像(HR)。SR的应用领域很广泛,如目标检测(尤其是针对小对象)、监控视频、卫星遥感等。实际上,无限多的高分辨率图像可以通过下采样得到相同的低分辨率图像。因此,SR问题本质上是不适定的,不存在唯一解。为了缓解这一问题,文献中提出了许多SISR方法,目前的SISR方法一般可分为两类:基于重建的方法(Reconstruction-Based)[1-2]和基于样本学习的方法(Example Learning-Based)。基于重建的方法仅用了一些先验知识来正则化重建过程,计算量低且简单,不过无法处理复杂图像结构。基于学习的方法利用大量通过学习得到的先验数据来学习低分辨率图像到高分辨图像的某种映射关系,用学习到的映射关系来预测高分辨率图像,常见的有基于稀疏编码和基于深度学习的方法。

近年来深度学习发展迅猛。Dong等[3]首先提出了一个采用三层卷积神经网络的SRCNN模型对特征提取、非线性映射和图像重建阶段进行端到端的联合优化。Shi等[4]提出了一种有效的亚像素卷积神经网络ESPCN,它直接在LR空间中提取特征映射,有效地降低了计算复杂度。受到功能强大的深度卷积神经网络(Convolutional Neural Networks,CNN)的影响,许多基于CNN的SR方法尝试训练一个更深的网络以获得更好的重建性能。Kim等[5]提出了一种包含20个卷积层的VDSR模型,该模型采用残差学习和自适应梯度裁剪来减轻训练难度,证明了更深的网络可以提供更好的特征重建,使网络捕捉到更加全局的特征,进而有益于超分辨率重建工作中图像细节的恢复。

因此在之后提出的一些模型[6-9]为了获得更好的性能,不断深化或拓宽网络已经成为了一种趋势。然而,不断对网络进行深化和拓宽,随之而来的则是实际应用中算法复杂度和内存消耗的提升。在实际应用中,如移动和嵌入式视觉应用中,这些方法的适用性较差。此外,传统的卷积层平等地对待图像的高频特征与低频特征,不加区分地将特征传向后续层。但高频特征往往包含更多的细节信息,而SR任务的重点是恢复图片中的纹理细节信息,因此应该将重点放在拥有高频细节的特征上。随着研究的深入,注意力的重要性得到了广泛的关注,注意力不仅阐明了关注的焦点在哪里,还提高了焦点的表达。与被广泛关注的Hu等[10]提出的只关注通道注意力的方法不同,Woo等[11]对基于通道和空间两个维度上的特性响应进行了研究,提出了通道和空间注意力机制。它可以自适应地重新校准通道和空间两个维度上的特性响应以提高网络的表达能力,很好地解决了深层网络中特征不加区分的向后传递的问题,同时该注意力机制引入的参数量很少。这为SR获取全局特征来更好地恢复高频纹理细节以及网络的轻量化提供了新思路。

基于上述讨论,本文提出一种基于双注意力机制的轻量级图像超分辨率重建(SR-LAM)算法。所提出的SR-LAM网络为了能够充分地捕获上下文信息,将网络深度保持在20层。网络结构分为四部分。第一部分由一个增强模块对LR图像进行浅层特征的提取及校准。第二部分由包含组卷积的卷积核大小为3×3的特征提取块堆叠,逐步细化残差信息。第三部分与第一部分相同,用一个增强模块进行深层特征提取和重新校准。最后由重建模块将残差信息上采样到目标尺寸后与上采样到目标尺寸的输入图像融合输出结果。最后通过在纵向和横向上进行对比实验,表明了尽管使用较少的卷积层,本网络仍然能够获得具有竞争性的结果。

1 方法与理论

自Kim等[5]的方法被提出以来,许多超分辨率重建算法都是基于它进一步优化网络结构以提高网络性能。本文也将在VDSR模型的基础上提出改进。

该模型以插值后的目标尺寸低分辨率图像作为网络的输入,这无疑加大了网络的计算量。为了拥有更大的感受野让图像特征提取到更多的细节,VDSR采用了更深(20层)的网络结构,为了加速收敛,采用了非常高的学习率,使用残差学习来缓解高学习率导致的梯度问题。Kim等通过实验证明了不断加深网络会带来感受野的增加,使网络能够捕捉到更加全局的特征,进而有益于SR工作中纹理细节的恢复。

虽然Kim等对超分辨重建算法做出了很好的改进,但其中仍存在着一些影响超分辨率重建质量和效率的因素:

1) 对于目标尺度系数n,与插值后的目标尺寸LR图像进行卷积的计算量将是与原始LR图像进行卷积计算的n2倍,另外,双三次插值方法不会带来额外的有利于解决不适定问题的信息。

2) VDSR及之后提出的一些模型为了获得更好的性能,不断深化或拓宽网络已经成为了一种趋势,但随之而来的是庞大的计算量及网络训练困难。

3) LR图像与HR图像之间存在许多相似的部分,SR任务的重点则是恢复图片中的纹理细节信息,因此应该将重点放在拥有高频细节的特征上。VDSR及一些传统的卷积网络认为所有空间位置和通道对超分辨率具有统一的重要性,通常采用级联网络拓扑,每一层的特征映射被无区别地发送到后续层,这不利于把注意力集中在更有利于恢复高频细节的特征上。

针对上述现象,随着研究的深入,本文对其做出了改进:

1) 本文方法采用后上采样代替提前上采样,以此来降低网络运算量,这有利于网络轻量化的实现。

2) 在某些情况下,考虑到并非所有特征都是SR所必需的,其重要性各不相同,有选择地侧重于某些特定特征对更好地恢复高频细节是有帮助的。由于卷积运算是通过混合跨通道和空间信息来提取特征的,所以强调沿着通道和空间轴这两个主要维度的重要特征是有必要的。对空间和通道两个维度上的特征响应研究[11]允许这种灵活性的产生。提出双注意力机制可以自适应地重新校准通道和空间两个维度上的特性响应,提高网络的表达能力,解决深层网络中特征不加区分的向后传递的问题。另外该注意力机制引入的参数量非常少。融合高效且轻量的双注意力机制代替普通卷积块在保证了网络质量的情况下进一步实现了轻量化,很好地解决了网络加深带来的训练困难。

3) 用组卷积代替普通卷积已经被证明[12]了能够在性能损失很小的前提下减少大量的网络参数,本文方法为了进一步实现网络轻量级在特征提取模块融入了组卷积。

1.1 模型结构

图1展示了本模型的总体网络架构。模型由四部分组成,包括三种模块:增强模块、特征提取模块、重建模块。

图1 SR-LAM网络结构

SR的具体任务是将原始高分辨率图像(IHR)进行下采样得到输入网络的低分辨率图像(ILR),经网络得到最终的重建后高分辨率图像(ISR)。其中的下采样操作其实是在模拟原始图像在其获取、存储和传输过程中存在不可避免的各种图像退化过程,退化函数通常是未知的,在目前的SR方法中大多采用双三次插值对IHR进行下采样预处理来模拟这个过程,为了与其他网络保持一致,本文也采用此插值方法来处理原始高分辨率图像。

对于网络的输入图像,第一部分由一个增强模块对LR图像进行浅层特征的提取及校准。其中增强模块包含一个卷积核大小为3×3的卷积层和基于通道和空间的双注意力模块。第二部分由5个包含组卷积的卷积核大小为3×3的特征提取块堆叠,逐步提取残差信息。第三部分与第一部分相同,用一个增强模块进行深层特征提取和重新校准。最后由重建模块将信息采用[4]提出的亚像素卷积的方式上采样到目标尺寸并和直接双三次插值上采样后的LR图像融合相加。

1.2 增强模块

增强模块由一个3×3的卷积和一个基于通道和空间两个维度的双注意力模块组成。卷积层后面跟着一个LReLU激活函数,如图2所示。首先进行浅层特征的提取,紧接着对特征进行一次校准,让随后的特征提取步骤把注意力集中在更有利于恢复高频细节的特征上。输出维度为64的特征图。用公式表示如下:

图2 增强模块

Fout=FS(FC(H(Fin)))

(1)

式中:Fin表示输入;H表示特征提取函数;FC表示通道注意力操作;FS表示空间注意力操作;Fout为增强模块的输出。

此注意力模块是一个轻量级模块,在仅引入了1 194个参数的情况下提高了网络的表达能力,满足了对网络轻量级的要求。该模块将注意力的特征细化应用于通道和空间两个不同维度,加强了对高频纹理细节的校准能力,引入该模型使得本网络可以在保持较小开销的同时获得有竞争力的重建效果。

在通道注意力子单元中,如图2所示。对于输入的特征图F∈RH×W×C首先使用平均池化和最大池化来聚合特征图的空间信息,生成两个不同的空间上下文信息来分别表示平均池化特征和最大池化特征。接着将两个上下文特征分别发送到共享网络,共享网络由含有一个隐藏层的MLP组成。之后对MLP输出的两个特征进行逐元素求和以及Sigmoid激活操作,生成通道注意力特征图AC(F),该特征图和输入特征图逐元素相乘得到空间注意力子模块的输入特征图FC,这个过程可以表示为:

FC=sigmoid(MLP(Avgpool(F))+MLP(Maxpool(F)))

(2)

在空间注意力子单元中,如图2所示。通道注意力子模块的输出Fc即本子模块的输入。首先对其沿着通道轴进行平均池化和最大池化得到两个上下文信息,将得到的两个信息基于通道方向做拼接,经一个卷积降维到单通道,然后经Sigmoid激活函数生成最终的空间注意力特征图。最后将该特征图与输入特征图做乘法得到FS,这个过程可以表示为:

FS=sigmoid(f7×7(Avgpool(FC);Macpool(FC))

(3)

式中:f7×7为一个7×7×1的卷积层。

上述过程可以简洁地概括为:

FC=AC(F)⊗F

(4)

Fout=FS=AS(FC)⊗FC

(5)

式中:AC为通道注意力操作;AS为空间注意力操作。

1.3 特征提取

用组卷积代替普通卷积已经被证明[12]可以减少大量的参数和操作,而性能损失很小。为了进一步降低网络参数,减少网络计算成本。本文在特征提取模块的第二层引入了组卷积。组卷积可看成是一种稀疏卷积连接的形式。将输入特征在通道方向上分成若干组,分别对每一组进行卷积后在通道方向进行拼接。如图3所示。特征提取模块由两个3×3的卷积层和一个组卷积块组成。使用组卷积的优势在于它使模型的效率可调,用户可适当地选择组大小,因为组大小和性能处于权衡关系中。本文延续Ahn等[12]所提出结论,设定组卷积的组大小为4。

图3 特征提取模块

1.4 重建模块

早期的上采样方法大多是采用双三次插值进行提前上采样,即在网络的第一层之前或者第一层对输入的LR上采样到与输出的HR相匹配的尺寸。这种方法显然会增加计算复杂度,尤其是对于卷积操作,因为卷积网络的处理速度直接取决于输入图像的分辨率,这不符合本文轻量级网络的目标。其次,双三次插值方法不会带来额外的有利于解决不适定问题的信息。因此本文没有延续VDSR中的提前上采样操作,而是采用了文献[4]中提出的亚像素卷积进行后上采样。亚像素卷积通过通道扩增和像素点重排来实现图像放大。由于输入的图片分辨率降低,可以有效地使用较小的卷积核来进行特征提取,同时维持给定的上下文区域。分辨率和卷积核尺寸的减小也大大降低了计算量和内存的复杂度。

最后将重建模块的输出与经上采样到目标尺寸的输入图像融合得到最终结果。

1.5 损失函数

本文采用峰值信噪比(PSNR)和结构化相似度(SSIM)作为重建的评价指标。PSNR与还包含L1和L2的像素损失高度相关,像素损失最小直接最大化PSNR。本文采用L1作为模型的损失函数,L1公式如下:

(6)

2 实验训练及结果

2.1 数据集

该模型使用DIV2K数据集[13]进行训练,它包括1 000幅图像,其中:800幅训练图像;100幅验证图像;100幅测试图像。它是一种新提出的高质量图像数据集。由于该数据集的丰富性,最近的SR模型[14-16]也使用了DIV2K数据集。本文方法在四个被广泛使用的标准数据集上进行评估:Set5、Set14、BSD100、Urban100。

2.2 实验细节

本文使用文献[17]中提出的Xavier初始化方法进行权重初始化,偏置项初始化为零。使用Adam[18]对网络进行优化。初始学习率为:1E-4,最大迭代次数为10 000,每2 000次迭代学习率减半。

本文的实验环境如表1所示。

表1 实验环境

2.3 实验结果及分析

为了探索通道和空间的注意力机制对SR任务的作用,首先将本文模型与剔除掉注意力模块的NA-SR-LAM进行了对比。训练过程如图4所示,横轴为训练轮数,纵轴为PSNR值。可以看到包含注意力模块的模型训练过程更稳定,随着训练次数的增加,波动逐渐变小。实验结果如表2所示。实验结果显示迭代次数为10 000时,添加了通道和空间注意力的模型较没有添加的模型在四个测试数据集上的峰值信噪比值分别提高了:0.29 dB,0.27 dB,0.24 dB,0.53 dB。此外,一个普通的3×3卷积核将引入36 928个参数,一个注意力模块仅引入了1 194个参数。综上所述,通道和空间注意力模块不仅有利于增强特征表达能力而且有利于模型的轻量化。

表2 SR-LAM与无注意力的SR-LAM对比

图4 SR-LAM与无注意力的SR-LAM对比(S=2)

表3展示了迭代次数为3 000时,Set14数据集上有无组卷积情况下模型参数量以及PSNR的对比。其中N-Group表示用普通卷积来代替特征提取模块中组卷积。

表3 SR-LAM与无组卷积的SR-LAM对比

由表3可知,融合组卷积的模型与未应用组卷积的模型在PSNR上仅相差0.01 dB,但参数量却比未应用组卷积的模型降低了将近23%,因此引入组卷积是实现该模型轻量化的有效措施。

本文提出的改进策略在不同程度上都对网络的轻量化做出了一定贡献,为了进一步直观展示本文算法是否更好地平衡了模型复杂度与重建效果,本文对近年来基于深度学习的经典超分辨率算法进行了对比实验。实验条件:目标尺度系数(S)为4,数据集为Set14。

选取的对比算法有:SRCNN[3]、LapSRN[19]、FSRCNN[20]、VDSR[5]、DRRN[6]、IDN[16]、本文提出的SR-LAM算法:SRCNN模型采用三层卷积神经网络对特征提取、非线性映射和图像重建阶段进行端到端的联合优化;FSRCNN方法将SRCNN的提前上采样改为了后上采样,并且用8个小尺寸的卷积核来代替之前的大卷积核,降低了计算复杂度,重建效果略有提升;VDSR将网络加深到20层以提高图像重建效果;DRRN延续了VDSR更深的网络层次重建效果更好的思想,将网络加深到52层,重建效果略有提升,模型参数量也有所下降;LapSRN方法对原始LR图像先生成低倍放大图像,再逐步细化生成高倍放大的图像,很好地改善了放大倍数高的复杂度问题;IDN直接从LR图像提取特征,减少了模型计算量,受到注意力思想的影响,作者提出了可提高网络表达能力、压缩特征冗余信息的蒸馏模块,通过对此模块的叠加在一定程度上提高了网络重建效果。

实验结果如图5所示,其中横轴表示算法所需参数量,纵轴为重建效果评价指标PSNR,不同大小的圆点代表不同算法所需参数量。结果表明:本文提出的方法更好地兼顾了重建效果与模型复杂性。与重建效果相当的IDN、LapSRN相比,本文所提出的模型更加简洁,参数量更少。在模型参数量方面,SR-LAM的重建效果远好于参数量相当的DRRN。对比基础模型VDSR,其参数量是本文所提方法的1.5倍。

图5 不同算法在Set14数据集上的PSNR及参数个数对比(S=4)

图6直观地对比了目标重建系数为2的comic图像的重建效果。对比Bicubic、SRCNN、VDSR、IDN方法,本文方法视觉上手指边缘以及流水的线条都更加清晰,这是因为本文将更多的注意力放在了高频信息,因此本方法对线条以及高频信息的重建效果更好、纹理更加清晰。

图6 不同方法对comic重建的视觉效果对比(S=2)

表4显示了最近基于CNN的SR方法的性能和模型大小。选取的算法包括:LapSRN[19]、IDN[16]、RCAN[21]、SAN[22]。本文提出的SR-LAM较LapSRN和IDN拥有更少的参数和更好的性能。虽然RCAN和SAN性能略胜,但其付出的代价特别大,二者平均参数量是本文所提方法的32倍。

表4 参数及性能比较(Set14)

表5客观地对比了本文方法与其他五种经典方法(双三次插值(Bicubic)、SRCNN、FSRCNN、VDSR、IDN)在不同上采样倍数下的PSNR以及SSIM。

表5 本文方法(SR-LAM)与其他方法结果比较

结果显示本文模型取得了具有竞争力的结果:与基础模型VDSR相比,在S分别等于2、3、4的情况下,四个测试数据集上的PSNR值平均提高了0.273 dB、0.253 dB、0.295 dB,SSIM值平均提高了0.003、0.005、0.008;与先进的IDN模型相比,在Set14和Urban100数据集上本文方法取得了更好的结果,而且通过对比图6(e)和图6(f)可以发现,本文方法对手指及指甲处的边缘恢复得更为清晰,视觉效果更好,与原图像也更加相似。

3 结 语

本文提出一个简洁、轻量、便于在移动端使用的超分辨率重建模型:SR-LAM。方法上主要通过:融合高效且轻量的双注意力机制代替普通卷积块、采用后上采样方法来代替提前上采样、引入部分组卷积的方式,使模型能够在拥有少量网络参数的情况下达到一个有竞争力的重建效果。实验结果证明,对比重建效果相当的先进模型,本文方法仅使用了其1/2,甚至更少的参数量。因此可以说本文模型更好地平衡了模型复杂度与重建效果。后续将进一步研究如何将这种高效、简洁的网络用于视频超分辨率领域。