多分支残差特征蒸馏网络的图像超分辨重建

李 轩,刘立柱

(沈阳航空航天大学 电子信息工程学院,沈阳 110136)

1 引 言

图像超分辨是从低分辨率(low-resolution,LR)图像中恢复出相应的高分辨率(high-resolution,HR)图像,它是图像处理和计算机视觉任务中的一个重要问题.由于硬件设备限制或者恶劣环境等原因,通常会获取到一些LR图像.图像超分辨率技术可以在不增加硬件成本的前提下提升图像质量,在安防监控、卫星图像、显微成像及医学影像等领域有广泛应用[1].超分辨率可以大致分为两大类:传统方法和基于深度学习的方法.近年来,得益于深度学习方法的强大学习能力,基于深度卷积神经网络的方法在重建超分辨率图像方面取得了很大进展.接下来简要回顾了与卷积神经网络相关的图像超分辨率工作.

Dong等人[2]首次提出一个三层卷积神经网络SRCNN,以端到端的方式进行图像重建,呈现出良好的重建效果;受这项工作启发,Kim等人基于残差学习的方法设计了20层的VDSR[3]以提高超分辨率(super-resolution,SR)性能,作者之后结合递归监督方法提出了16层递归的DRCN[4];Tai等人对DRCN做出改进,基于局部残差连接和残差单元的递归学习方法提出了DRRN[5],获得了更好的性能,之后又提出MemNet[6],该网络引入了一个包含门控单元和递归单元的内存块,以通过自适应学习过程显式地挖掘持续记忆;Lai等人[7]提出了拉普拉斯金字塔超分辨率网络(LapSRN),将LR图像作为输入,逐级放大,不仅减少了计算量,还提高了重建性能;为了减少计算量和提高速度,Shi等人[8]在ESPCN中设计了一种有效的亚像素卷积层,用来在网络末端进行LR到SR的特征映射,使得大部分计算在低维特征空间中进行;Dong等人[9]提出了快速自适应卷积神经网络(FSRCNN),它采用转置卷积作为上采样层生成SR图像;Lim等人[10]提出了EDSR和MDSR,通过去除残差网络中的批量标准化(Batch Normalization,BN)层显著提高了性能;Zhang[11]等人通过在残差块中引入密集连接,提出了RDN,进一步提高了性能.

在最新研究中,注意机制在计算机视觉任务中被广泛使用,Hu等人[12]提出了压缩提取网络(SENet),实现了图像分类的显著改进.最近一些学者将注意力机制引入超分任务中用于提升性能,Zhang等人[13]提出了残差通道注意力网络(RCAN),通过注意力机制聚焦于边缘纹理特征;Dai等人[14]提出了一个二阶注意力网络(SAN),具有更强大的特征表达和特征关系学习能力;Liu等人[15]提出一个残差特征融合网络(RFANet),将连续残差块特征融合,并提出一种空间注意力模块,结合两个模块,最终模型的超分速度和精度都优于RCAN.尽管上述基于CNN的方法取得了巨大的成功,但大多数方法参数量巨大.为了解决这个问题,Tian等人[16]提出非对称式卷积神经网络(ACNet),利用非对称式卷积进行特征提取,对高低频特征分别进行上采样并融合得到SR图像;Ahn等人[17]通过级联网络架构提出了移动场景的CARN-M模型;Hui等人[18]提出了信息蒸馏网络(IDN),该网络采用通道分离的方式提取蒸馏特征,在IDN的基础上,还提出了快速和轻量级的信息多重蒸馏网络(IMDN)[19].

虽然SISR已经取得了相当大的进步,但现有的基于CNN的模型仍然面临一些限制.模型计算量大很难在移动端部署;传统的网络结构仅仅通过3×3的卷积核提取特征,这忽略了局部像素点的影响,且存在大量平坦区域的特征提取,导致训练效率低;现有的超分模型几乎都只是将残差学习作为一种策略来缓解训练难度,没有充分利用残差特征,然而残差特征对于超分至关重要.本文针对这些问题,提出一种多分支残差特征蒸馏算法,主要贡献包含以下几点:

1)在特征提取阶段,提出一种多分支残差模块,通过多分支结构更好地利用多尺度信息和上下文信息,在此基础上增加了一个通道注意力模块,获取不同特征层之间的关联信息;同时添加了残差连接,利用多分支卷积、注意力机制及残差学习获得更加精细的特征.

2)在特征蒸馏和融合阶段,结合1×1卷积和Meta-Acon(Meta activate or not)自适应激活函数进行降维,减少平坦区域的大量冗余计算,从而保证性能的同时降低模型复杂度.提出瓶颈注意力模块,该模块结合多分支空洞卷积,构建出更有效的空间注意力图,从而更好地提取出边缘、纹理等特征.

3)与现有轻量化模型相比,本文所提算法模型有更好的超分重建结果,评价指标PSNR和SSIM进一步提升;在2倍超分实验中,超分结果已经达到部分大型模型的水准.

2 多分支残差特征蒸馏网络

为了在SR性能和模型复杂度之间达到一个较好的平衡,为此本文提出了多分支残差特征蒸馏网络(图1).它主要包含3个部分,特征提取、特征蒸馏及特征融合、重建.

图1 多分支残差特征蒸馏网络结构Fig.1 Residual feature distillation network via diverse branch convolution

本节将具体介绍残差特征蒸馏模块(Residual feature distillation block,RFDB)、多分支残差模块(Diverse branch residual block,DBRB)、瓶颈注意力模块(Bottleneck Attention Module,BAM)及损失函数.

2.1 残差特征蒸馏模块(RFDB)

Hui等人[19]提出了IMDN网络,该网络采用通道分离的方式提取蒸馏特征,但是这种方式不灵活,且采用3×3卷积会生成很多冗余参数.故本模块利用1×1卷积进行降维,其等价于通道分离操作,且可以更好地利用通道间的关联信息;利用多分支残差模块(见2.2节)进行深度特征的提取;在局部残差特征融合阶段,利用瓶颈注意力模块(见2.3节)聚焦于边缘纹理特征.三者组成了残差特征蒸馏模块,如图2所示.

图2 残差特征蒸馏模块Fig.2 Residual feature distillation block

对于图2中的特征蒸馏操作是通过卷积以一定比例压缩特征通道实现,采用1×1的卷积进行通道降维,不仅高效,还可以去除大量冗余参数,进一步降低参数量.在完成逐级特征蒸馏计算后,利用自适应激活函数和1×1卷积对所有蒸馏特征进行去冗余和降维、融合操作,然后使用瓶颈注意力模块聚焦于边缘、纹理特征,最后利用全局残差连接对输入特征进行增强.

设输入图像为x,特征提取阶段可以描述为:

F0=h0(x)

(1)

其中h0表示3×3卷积,F0是提取的特征.

在残差特征蒸馏阶段,将低频特征F0作为输入,进一步提取特征,该过程为:

Z0=H0(F0)

(2)

其中Z0表示提取到的特征,H0代表残差特征蒸馏模块,然后通过多个以链式方法堆叠的残差特征蒸馏模块进一步提取特征,该过程为：

Zk=Hk(Zk-1),k=1,2,…,n

(3)

其中Hk表示第k+1个残差特征蒸馏模块,Zk-1表示输入特征,Zk表示输出特征.在完成逐级特征计算后,然后利用1×1卷积对所有中间特征进行降维、融合,再利用自适应激活函数调整中间特征图,减少冗余参数,最后通过3×3卷积进一步提取特征.该过程为：

(4)

其中F表示融合特征,Acon是自适应激活函数,f11×1是1×1卷积,f23×3是3×3卷积.

Meta-Acon是Ningning等人[20]提出的,可以自适应的选择是否激活神经元.计算公式为:

y=(p1x-p2x)*σ(β*(p1x-p2x))+p2x

(5)

其中x是输入特征图,σ是sigmoid激活函数,p1和p2是两个可学习参数,β是一个自适应函数,用来控制是否激活神经元.β的计算公式为：

β=σ(fc(fc(GAP(x))))

(6)

其中fc表示1×1的卷积,GAP表示全局平均池化.

低分辨率图像的细节主要在边缘、纹理区域,平坦区域缺失较少,但特征提取存在大量平坦区域的冗余计算,通过该自适应激活函数,可以减少冗余参数,结合1×1卷积降维,充分利用通道间信息,可以在保证性能的同时降低模型复杂度.

2.2 多分支残差模块(DBRB)

Ding等人[21]提出非对称式卷积,利用非对称编码思想,采用两个一维卷积从水平和垂直方向对3×3卷积进行增强;Jia等人[22]利用拉普拉斯金字塔结构渐进重建SR图像;Cheng等人[23]利用交叉学习机制,进一步提升SR重建性能.本文从提取多尺度特征的角度出发,构建了多分支卷积模块,用于对3×3卷积进行增强.

多分支卷积模块(图 3)利用空洞率分别为2、4、6及8的多个空洞卷积分支,可以更好地结合多尺度信息和上下文信息.这种多分支结构可以更好地提取边缘、纹理特征,丰富3×3卷积得到的特征,从而提高模型的特征提取能力.

多分支残差模块(图 3)在多分支卷积模块的基础上增加了一个通道注意力模块,充分利用通道之间的关系,获取不同特征之间的关联信息;同时增加了一个恒等连接分支,可以在不引入额外参数的前提下从残差学习中受益,从而充分利用注意力机制、多分支卷积及残差学习提取精细特征.

图3 多分支残差模块Fig.3 Diverse branch residual block

2.3 瓶颈注意力模块(BAM)

注意力机制会逐渐聚焦于高层次特征,类似于人类的感知机制,近年来在图像分类、图像分割及目标检测等领域广泛应用,常见的注意力机制:通道注意力和空间注意力.注意力机制会让网络更加关注对任务有益特征,抑制无用的特征(可以理解为在神经网络隐藏层加权注意力机制),从而更合理分配地计算资源[24].

正如在引言中提到的,在计算机视觉领域,Hu等人提出的SENet,利用全局平均池化特征来学习不同通道特征的重要程度,在图像分类上任务上有显著改进;Woo等人[25]融合通道注意力和空间注意力提出Convolutional Block AttentionModule(CBAM),与SENet不同,它在通道和空间上同时做全局最大池化和全局平均池化,用来获取更多信息,该模块相对于SENet多关注了空间位置特征,进一步提高了模型的特征提取能力.通过引言介绍可知,RCAN、SAN和RFANet等将注意力机制引入超分领域,都显著提高了超分性能.

故本文提出了瓶颈注意力模块,如图4所示,其会增强或抑制不同空间位置的特征,从而消除像背景语义特征这样的低层次特征,然后逐渐聚焦于边缘、纹理等高层次特征.因为一个大的感受野可以充分利用上下文信息获取高层次特征,所以采用空洞卷积来有效的扩大感受野.结合多分支空洞卷积,充分利用上下文信息和多尺度信息,从而有助于构建更有效的空间注意力图.

图4 瓶颈注意力模块Fig.4 Bottleneck attention module

空洞卷积:扩大感受野和捕获多尺度上下文信息.池化等降采样操作可以增大感受野,但是降低了空间分辨率.而空洞卷积可以在不丢失分辨率的情况下,增大感受野,从而更好地提取高层次特征.空洞卷积的空洞率不同,感受野就不同,也就获取到了多尺度信息.因为空洞卷积的相邻像素是通过不相关的子集卷积得到的,缺少相关性.如果使用相同的空洞率,就会导致明显的网格效应.本文结合分解卷积的思想,首先通过两次1×1卷积获取局部信息,然后采用2、4、6、8的分层空洞率进行卷积,不用一个固定的空洞率,再结合后续的1×1卷积,可以缓解网格效应问题,且大大增加了感受野,从而更好地获取上下文信息.具体过程如下:

首先利用1×1的卷积进行降维,然后使用4个不同空洞率的空洞卷积来有效利用上下文信息和多尺度信息,将拼接后的特征与残差相加,其次通过1×1卷积增加通道间的关联,再用3×3的卷积进一步提取特征,最后使用1×1的卷积升维得到空间注意力图Y(F)∈C*H*W.空间注意力图的计算公式为:

(7)

(8)

(9)

(10)

其中f11×1、f21×1、f71×1和f91×1表示1×1的卷积,f33×3、f43×3、f53×3和f63×3表示卷积核为3×3、空洞率分别为2、4、6和8的空洞卷积,σ是sigmoid激活函数.

对于给定的特征图F∈RC*H*W,该模块可以得到一个注意力图Y(F)∈RC*H*W,聚焦后的特征图为:

F′=F*Y(F)

(11)

其中*表示元素相乘.

2.4 损失函数

本文采用L1损失函数来评估重建的超分辨率图像和真实高分辨率图像的差异,进而调整网络参数来优化网络.

损失函数可以表示为:

(12)

3 实验和分析

3.1 数据集和评价指标

现有方法大都是使用DIV2K数据集作为训练数据集,为了实验公平,本文也选择该数据集来训练超分模型.该数据集由3部分组成,800张训练图片,100张验证图片和100张测试图片.因为该数据集是2K分辨率的,在训练时往往不需要这么大,预处理阶段将其裁剪为有重叠的480×480的子图像块存储在硬盘中;在MATLAB中使用双三次线性插值对高分辨率图像分别进行2、3、4倍的下采样得到低分辨率图像,同理对于2、3、4倍的下采样图片分别裁剪出240×240、160×160、120×120的子图像块.

为了客观评估本文所提网络的超分性能,在测试阶段,使用广泛使用的基准数据集Set5,Set14,BSD100与当前一些先进的超分模型进行对比.采用峰值信噪比(peak signal-to-noise ratio,PSNR)和结构相似性(structural similarity,SSIM)去评估超分辨率图像的质量.所有值是在Ycbcr颜色空间的Y通道上进行计算得到的.

3.2 训练配置和参数细节

网络运行在Pytorch框架下,Ubuntu系统18.04,显卡为NVIDIA Tesla V100 GPU(32G).在训练阶段从低分辨率子图像块中随机裁剪出64×64的图像块作为输入,采用随机水平翻转和旋转来进一步扩充训练数据集.采用Adam优化器训练模型,参数为β1=0.9,β2=0.999,ε=1×10-8,初始学习率为5×10-4,Batch size为64,训练过程一共80万次迭代,每20万次迭代学习率减小一半.

3.3 模型分析

为了验证本文方法的有效性,用3×3卷积代替多分支残差模块(DBRB),用ESA和CBAM模块代替瓶颈注意力模块(BAM),去掉Meta-Acon模块.其中,ESA是在RFANet提出的增强型空间注意力模块,CBAM是前面提到的一种混合注意力模块.

由表1可以看出,“3×3卷积+BAM”的组合与“3×3卷积+ESA”和“3×3卷积+CBAM”这两种组合策略相比,参数量相差不大,在基准数据集上的2倍超分结果上,PSNR和SSIM均有提高,充分证明了BAM可以充分利用多尺度信息和上下文信息提取特征;“DBRB+ESA”的组合与“3×3卷积+ESA”的组合相比,参数量增加了152K,在2倍超分结果上,PSNR和SSIM有明显提高,PSNR平均增加0.11dB,SSIM平均增加0.0014,这表明DBRB在适当增加计算量的前提下,更好地提取到精细特征;“Acon+DBRB+ESA”的组合与“DBRB+ESA”的组合相比,参数量增加了2K,在2倍超分结果上,PSNR和SSIM进一步提高,PSNR平均增加0.03dB,SSIM平均增加0.0004,这证明了Acon可以很好地关注“感兴趣”特征,提高模型表达能力;“Acon+DBRB+BAM”的组合与“Acon+DBRB+ESA”的组合相比,参数量减少了15K,在2倍超分结果上,PSNR和SSIM进一步提高,PSNR平均增加0.07dB,SSIM平均增加0.0007.

表1 不同模块对网络的影响Table 1 Influence of different modules on the network

通过分析,证明了本文所提多分支残差模块、瓶颈注意力模块在提取特征上有效性,并且证明了所引用的自适应激活函数的有效性,它们的组合可以获得更好的超分性能.

3.4 方法对比

为了验证本文算法的有效性,将其与16个超分算法进行了对比,包括VDSR、DRCN、DRRN、MemNet、LapSRN、EDSR、RDN、RCAN、SAN、RFANet、ACNet、CARN、IDN、IMDN、SRMDNF[26]、DBPN[27].表2展示了不同算法在基准数据集上分别进行2、3、4倍放大后的平均PSNR/SSIM.

表2 不同方法在基准数据集上分别进行2、3、4倍上采样后的平均PSNR/SSIMTable 2 Average PSNR/SSIM for scale factor 2,3 and 4 on benchmark datasets

与轻量化模型相比:由表2可知,在前10行的轻量化模型中,IMDN的PSNR最高.从表3可知,本文算法与IMDN相比,在Set5、Set14、BSD100和Urban100数据集上的2倍超分结果上,PSNR平均提高0.16dB,SSIM平均提高0.0019;3倍超分结果上,PSNR平均提高0.15dB,SSIM平均提高0.0027;4倍超分结果上,PSNR平均提高0.12dB,SSIM平均提高0.0044.从定量角度证明了本文算法较现有轻量化超分算法有了一定提高.

表3 本文方法与IMDN的对比Table 3 Comparison of IMDN and our method

与大型模型相比:由表2可知,EDSR、DBPN、RDN、RCAN、SAN及RFANet的参数量都在千万级别,而本文算法,参数量只有70万.本文选取EDSR和DBPN来对比,由表4可知,从Set5、Set14、BSD100的2倍超分结果来看,DBPN的PSNR分别是38.09dB、33.85dB和32.27dB,本文算法的PSNR分别是38.09dB、33.82dB和32.25dB,PSNR基本相同,SSIM平均提高0.001;与EDSR相比,PSNR分别差0.02dB、0.1dB和0.07dB,SSIM基本相同.但是在Urban100的2倍超分结果上,本文方法较EDSR差距较大, PSNR差了0.48dB,SSIM差了0.0039,说明本文方法在城市这种大范围景观下的特征提取能力还有待提高.与新方法RCAN、SAN及RFANet相比,差距正在逐渐减小,而本文方法参数量却远远少于它们.从定量层面证明了本文算法在超分重建上的有效性.

表4 本文方法与DBPN及EDSR的对比Table 4 Comparison of the method in this paper with DBPN and EDSR

图5展示了在Set5数据集上进行4倍上采样的不同超分方法的参数量和性能对比图,从图中可以直观地看出,本文所提算法模型与轻量化模型IMDN、IDN、CARN和ACNet相比,有更高的PSNR;与大型模型DBPN和RDN相比,在性能接近的情况下,参数量分别只有其1/15和1/30.证明了本文算法在超分重建和特征蒸馏上的有效性.

图5 在Set5数据集上进行4倍上采样的参数量与PSNR的关系图Fig.5 PSNR vs.Parameters on Set5(X4)

除定量分析外,实验从DIV2K验证集随机选取若干下采样图片测试结果,从视觉层面对比不同超分算法的性能.如图6所示,为了更好地呈现超分图像细节对比结果,通过对图像的局部区域进行放大处理,来对比不同算法的超分性能.从图中可以看出,与轻量化算法ACNet、IMDN相比,本文算法在图像细节上有更好的重建效果;与大型算法RCAN相比,虽然还存在极少量细节模糊,但是已经基本重建出细节轮廓.

图6 在BI退化数据集进行4倍上采样效果对比图Fig.6 Visual comparison for scale factor 4 on Bicubic(BI)degradation datasets

4 结 论

为在有限资源的情况下获得更好的超分性能,本文提出了一种多分支残差特征蒸馏算法.在特征提取上,该算法提出的多分支残差模块,可以充分利用通道间信息、上下文信息和多尺度信息,提取更精细的特征;在特征蒸馏及融合阶段,利用卷积和通道自适应激活函数进行降维,保证性能的同时减少冗余参数,结合瓶颈注意力模块,进一步聚焦于边缘、纹理等特征,从而更好地重建出超分辨图像.实验结果表明,本文算法相对于现有轻量化模型有更好的细节特征提取能力,不仅有更高的PSNR和SSIM,还有较好的视觉效果;相对于大型模型,以1/15的参数量达到相近的重建效果.从定量和可视化层面证明了本文算法在模型性能和复杂度上达成了更好的平衡.