林承浩,吴丽君
(福州大学 物理与信息工程学院,福建 福州 350108)
图像超分辨率(Super Resolution,SR)是一项被广泛关注的计算机视觉任务,其目的是从低分辨率(Low Resolution,LR)图像中重建出高质量的高分辨率(High Resolution,HR)图像[1]。由于建出高质量的高分辨率图像具有不适定的性质,因此极具挑战性[2]。随着深度学习等新兴技术的崛起,许多基于卷积神经网络(CNN)的方法被引入到图像超分任务中[3-6]。SRCNN[3]首次将卷积神经网络引入到图像超分任务中,用卷积神经网络来学习图像的特征表示,并通过卷积层的堆叠来逐步提取更高级别的特征,使得重建出的图像具有较高的质量。在后续研究中,Kaiming He等人提出了残差结构ResNet[5],通过引入跳跃连接,允许梯度能够跨越层进行传播,有助于减轻梯度消失的问题,使得模型在较深的网络情况下仍然能保持较好的性能。Bee Lim等人在EDSR[6]中也引入了残差结构,EDSR实际上是SRResnet[7]的改进版,去除了传统残差网络中的BN层,在节省下来的空间中扩展模型尺寸来增强表现力。RCAN[8]中提出了一种基于Residual in Residual结构(RIR)和通道注意力机制(CA)的深度残差网络。虽然这些模型在当时取得了较好的效果,但本质上都是基于CNN网络的模型,网络中卷积核的大小会限制可以检测的空间范围,导致无法捕捉到长距离的依赖关系,意味着它们只能提取到局部特征,无法获取全局的信息,不利于纹理细节的恢复,使得图像重建的效果不佳[5]。
由于Transformer在自然语言处理(Natural Language Processing,NLP)领域中取得了较好的成效[9],Alexey Dosovitskiy等人将Transformer引入到计算机视觉(CV)领域中,即ViT(Vision Transformer),并且在多个视觉任务中取得了成功[10]。ViT的优势在于Transformer对于全局的信息更加敏感,模型中的注意力模块能够在输入序列的所有位置上进行全局交互,从而捕捉到长距离的依赖关系。但是,用于图像超分辨率的ViT,需要将输入图像分割成固定大小的块,并对每个小块进行独立的处理[11],这种处理策略就会产生两个弊端:一是修复后的图像可能会在每个小块周围引入边界伪影;二是每个补丁的边界像素可能会丢失信息,影响到重建图像的质量。Ze Liu等人提出了Swin-Transformer[11],将滑动窗口机制引入到Transformer中,其中滑窗操作包括了不重叠的Local-window和重叠的Cross-window,将注意力计算限制在一个窗口内可以大幅度节省计算量,并且通过滑窗操作也能使得注意力机制能够注意到全局的特征。
Jingyun Liang等人结合卷积神经网络和Swin-Transformer两者的优点提出了SwinIR[12],将两个构架以串行的方式应用在超分领域,展示出了混合构架在超分任务中的巨大应用前景。Peng等人提出了一种以CNN-Transformer并行的方法Conformer[13],它由一个CNN分支和一个Transformer分支组成,依靠特征耦合单元(Feature Coupling Unit,FCU)以交互的方式在不同分辨率下融合局部特征和全局特征,其结果表明了并行结构能以最大限度地保留局部特征和全局特征。虽然Conformer结合了两个网络的优势,但训练的参数量和训练时长也相应地增加了。并且,图像超分任务通常需要输入较高分辨率的图像,占用大量GPU内存,限制了模型的灵活性。要想取得更加高清的图像势必会增加网络模型中的参数量[14]。因此,在结合两种构架优势的同时,降低训练成本,使得图像超分辨率模型轻量化成为了本文需要解决的重大问题。
针对上述问题,本文提出了基于Swin-Transformer的单图像超分网络STSR(Swin-Transformer-based Single Image Super-Resolution)。具体贡献如下:结合CNN和Transformer的优势,设计了并行特征提取的特征增强模块(Feature Enhancement Block,FEB),能够有效地捕捉图像的局部细节特征,同时也能够捕捉长距离的依赖关系,使模型具有全局上下文建模的能力。本文采用了轻量化的Transformer模块,在达到较好重建效果的同时还能保持较低的计算成本。此外,通过设计动态调整模块(Dynamic Adjustment,DA),可根据输入图像的特征对输出进行动态的调整,从而增强网络的拟合能力,使得重建图像的纹理细节更加贴近于真实图像。
本文设计的基于Swin-Transformer的单图像超分网络STSR的整体网络结构如图1所示。结合CNN和Transformer两者的优势,本文设计了一个并行特征提取的特征增强模块FEB,采用卷积神经网络和Transformer网络并行提取图像特征,为了减少训练中的参数量,FEB中的轻量化Transformer模块只采用Swin-Transformer中的STL(Swin Transformer Layer)来提取特征,并且在网络中堆叠3层FEB以获取更深层次的特征,有助于提升重建图像的质量。为了进一步的提升重建图像的质量,本文设计了一个动态调整模块DA,动态调整模块主要采用的是通道注意力机制,通过这一机制,网络可以根据输入图像的内容和结构,自适应地增强或者减弱不同通道的影响,也可以减少网络对无关信息的依赖,使得在训练过程中网络能够快速收敛,减少训练时间,同时能够使得重建图像更加贴近于真实图像。STSR主要包括了四个模块:浅层特征提取、深层特征提取,动态调整模块和高质量图像重建模块。
图1 STSR的整体网络框架
(1)浅层特征提取
首先,使用一个3×3的卷积层来提取图像的浅层特征F0,如式(1)所示。
F0=H3×1(ILR)
(1)
(2)深层特征提取
因为浅层特征提取模块只经过一次简单的卷积操作,提取到较为初级的特征。将提取到的浅层特征F0,使用深层特征提取模块进一步地提取图像特征。深层特征提取模块由K个FEB模块和一个3×3的卷积层构成,用来提取深层特征FDF,如式(2)所示。
FDF=HDF(F0)
(2)
其中每个FEB的输出F1、F2、FK以及输出的深层特征FDF,如式(3)(4)所示。式中HKFi表示第i个FEB模块,Hconv表示最终的卷积层。
Fi=HKFi(Fi-1),i=1,2,…,K
(3)
FDF=Hconv(FK)
(4)
(3)动态调整模块
动态调整模块主要采用通道注意力机制,通过计算通道间的相关性和重要性,网络可以自适应地增强或减弱不同通道的影响。这有助于提取输入图像中最相关和有用的特征,使得重建图像与输入图像具有更高的相似度,从而改善超分辨率重建的效果。
动态调整模块接收来自浅层特征提取模块输出的浅层特征F0,经过处理后得到动态调整特征FDA,如式(5)所示。
FDA=HDA(F0)
(5)
(4)高质量图像重建模块
图像重建模块其实就是卷积与上采样的组合,本文中采用的是“卷积+pixel shuffle+卷积”的方式来进行图像重建。作为网络的最后部分,将接收浅层特征F0,深层特征FDF和动态调整特征FDA,以获得重建图像ISR。f和fp分别代表的是卷积层和亚像素卷积层,计算ISR的公式如下:
ISR=f(fp(f(FDF)))+f(fp(f(FDA)))+f(fp(F0))
(6)
本文方法中提出的FEB模块主要是由卷积神经网络CNN、轻量型Transformer和特征融合模块三个部分组成。特征增强模块FEB的大致结构如图2(a)所示。
图2 FEB模块和STL模块结构
(1)卷积神经网络
首先,利用卷积层来对浅层提取的特征F0进行进一步的特征提取,得到卷积层提取的特征FC,如式(7)所示。
FC=fc(F0)
(7)
其中fc为FEB模块中卷积层中特征提取的映射关系。
(2)轻量型Transformer网络
轻量型Transformer是由2个STL(Swin Transformer Layer)模块(如图2(b)所示)串联组成的,第一个STL模块中MSA(Multi-headed Self-Attention)采用的是W-MSA(Window based Multi-headed Self-Attention),W-MSA是窗口化的多头自注意力机制,相较于传统全局注意力机制减少了计算量。第二个STL模块采用的则是SW-MSA(Shifted Window based Multi-headed Self-Attention),由于W-MSA只能关注窗口本身的内容,无法跨窗口连接,这就导致了窗口之间的特征信息无法传递,而SW-MSA可以通过引入移位窗口的方式,在保持窗口化的情况下兼顾了全局特征,并提高了计算效率,使得模型达到轻量化的效果。使用多层感知机(Multi-Layer Perception,MLP),其中包括了两个全连接层和GELU激活函数,来进行进一步的特征变化。在MSA和MLP之前都添加了LN(LayerNorm)层,并且这两个模块都引入了残差连接。最后得到轻量化Transformer层提取的特征FT,如公式(8)所示。
FT=HSwin(F0)
(8)
其中HSwin为FEB模块中的轻量型Transformer层。
(3)特征融合模块
特征融合模块含有多尺度卷积块和特征重建部分,多尺度卷积块利用不同尺寸的卷积核来进行特征的多尺度提取,获得不同感受野的特征。特征融合模块采用残差连接的方式,将多个维度的特征相加,从而使网络更容易学习到低频和高频细节之间的映射关系,有助于加速训练过程。
在这个模块中,将卷积层提取的特征FC和轻量型Transformer层提取的特征FT来进行特征融合,得到FEB模块的输出特征FCf,如式(9)所示。
FCf=HCf(FC+FT)
(9)
其中HCf为特征融合模块的函数。
图3所展示的是动态调整模块,主要运用到的是全局平均池化(Global Average Pooling,GAP)操作。用于将卷积神经网络的特征图转换为固定长度的向量表示。GAP通过将每个通道的特征图转换为一个标量值。具体来说,假设输入图像经过卷积层后得到的特征图为F,其尺寸为H×W×C,其中H和W是特征图的高度和宽度,C是通道数,GAP会对每个通道的特征图计算平均值,得到一个C维的向量表示,记为向量A,即每个通道的平均值。总而言之,通过GAP操作,特征图的空间信息被压缩成一个固定长度的向量。对向量A进行一系列的线性变换和非线性激活操作,得到一个长度为C的通道权重向量S。将特征图F与通道权重向量S进行逐通道乘法操作,得到调整后的特征图F′。最后再将调整后的特征图F′输入到后续的网络层中,继续进行超分辨率重建的处理。
图3 动态调整模块作用流程图
通过动态调整模块,网络可以根据输入图像的内容和结构,自适应地增强或者减弱不同通道的影响,以提取到输入图像中最为主要的特征,同时减少网络对无关或者冗余信息的依赖,从而改善超分辨率重建的质量和细节保留的能力。
网络训练采用的损失函数是平均绝对误差( Mean Absolute Error,MAE),也称为L1损失函数,如式(10)所示。
(10)
在训练阶段,采用DIV2K数据集来训练模型。该数据集包含有1 000张高清图像(2K分辨率),其中包括800张图像作为训练数据,100张图像作为验证数据,100张图像作为测试数据。在评估阶段,采用的基准数据集是:Set5、Set14、BSD100、Urban100和Manga109。
在评价标准方面,采用峰值信噪比(PSNR)和结构相似度(SSIM)来作为衡量模型效果的评价指标。PSNR可以评价两幅图像之间的相似程度。PSNR值越高,说明重建出来的图像中的失真或者误差越小。较高的PSNR值表示重建图像质量较好,而数值较低意味着重建图像存在着更明显的伪影。SSIM与专注于对比像素差异的PSNR不同,SSIM主要比较图像的亮度、对比度和结构。通过计算三项的平均值来衡量相似性:亮度相似性、对比度相似性和结构相似性。所以SSIM更加侧重于图像的结构信息和感知质量。
网络训练所用平台为Ubuntu20.04,所有的实验均在单张 NVIDIA GeForce GTX 3060Ti显卡上完成训练。在训练之前,对HR图像进行不同缩放因子的双三次下采样,生成对应的LR图像。本文利用L1损失和Adam算法对模型进行优化。初始学习率定义为1×10-4,算法中参数设置为β1=0.9、β2=0.999,训练周期定义为300K iterations,并且采用余弦退火衰减的学习方案来加速收敛,动量momentum为0.9。
本文将STSR网络与一些经典的图像超分网络和近年来的轻量型图像超分网络进行对比,其中包括FSRCNN[4]、VDSR[15]、MemNet[16]、EDSR[6]、CARN[8]、IMDN[17]、LatticeNet[18]、ESRT[19]。表1为在同一实验环境下所获得测试结果。在每一行中,最好的结果用加粗的方式突出显示。从测试结果可以看出在模型较小的情况下,量化指标也取得了较好的结果,在模型性能和计算成本之间取得了较好的平衡。
表1 图像超分辨率重建效果量化比较结果
为了验证特征增强模块FEB以及动态调整模块DA的实际作用,本文将FEB模块和DA模块分别单独引入到STSR模型中,通过实验来验证这两个模块在量化指标PSNR和SSIM上的效果,可以验证出不同模块对网络性能提升的有效性。方法STSR-F是将STSR模型中的特征增强模块FEB替换为单一的卷积神经网络结构,由此来验证FEB模块的作用。方法STSR-D是去掉STSR模型中的动态调整模块DA,由此来验证DA模块在网络中的作用。
可以从表2中的量化指标得出特征增强模块FEB以及动态调整模块DA对于本文网络性能提升的有效性。
表2 消融实验
图4中给出了STSR与其他图像超分模型的视觉效果对比。本文设计模型重建出的高分辨率图像包含了更多纹理细节,对于边缘和线条等细节方面的重建,也展现出了更高质量的重建效果。
图4 STSR与其他模型的视觉对比
针对基于混合构架的图像超分模型较高计算成本的问题,本文提出了一种基于卷积神经网络CNN和Transformer混合模型,结合了两种构架的优势,提高模型对局部细节特征和全局信息的建模能力,增强了上下文信息的利用效率,并且通过轻量化Transformer网络,在达到较好重建效果的同时保持较低的计算成本。此外,本文还设计了DA模块,根据输入图像来动态调整网络的输出,使重建的图像更加贴近于真实图像。量化指标表明,本文方法在减少模型参数量的情况下,依然取得了较好的图像重建效果。