高效单图像超分辨率重建:深监督对称蒸馏网络

毛盼娣,徐道连

(1.重庆城市科技学院 电气工程与智能制造学院,重庆 402167;2.重庆大学 光电工程学院,重庆 400044)

单图像超分辨率(SISR)旨在利用低分辨率(Low-Resolution,LR)图像重建出对应的高分辨率(HR)图像,近年来该方法在物体检测[1]、信息取证[2]等领域应用十分广泛.DONG等[3]提出SRCNN首次为SISR引入了卷积神经网络(CNNs).随后,基于CNNs的SISR方法致力于设计更深更广的网络架构以实现性能提高.然而这类方法往往参数量和计算量巨大,例如EDSR[4]和RCAN[5]等.在现实场景下真正需求的是轻量化网络模型,尤其是可以部署到移动设备的实时模型,因而如何更好地实现网络轻量化是一个亟待解决的问题.

SISR模型轻量化首先想到的便是设计特殊的模型,减少网络的层数、使用小的卷积核、采用轻量化组件,例如分组卷积[6]和可变形卷积[7]等.但是特殊设计的模型依旧存在问题:一是轻量级网络与EDSR和RCAN等大网络相比具有很大的差距;二是轻量级模型中一些组件在实际硬件芯片上不能被很好地支持.SISR领域的蒸馏方法应运而生,其使用性能好的大模型(教师网络)蒸馏信息提升小模型(学生网络),使学生网络也能获得接近教师网络的性能.本文基于第二种方法探索.

SISR领域的蒸馏方法中,首次提出的SRKD[6]对教师网络和学生网络的不同阶段进行蒸馏,考虑到教师网络和学生网络通道数不相同,SRKD对中间信息的统计进行监督.随后,FAKD在SRKD的基础上引入了特征关联机制[8],进一步提升了学生网络的性能.但这些方法仍有一定的局限性,于是PISR首次引入特权信息(HR图像)[9],用教师网络提取HR高频特征补充到学生网络,明显提升了重建性能.利用特权信息的方法最近得到了明显的关注,基于参考图像的超分辨率方法C2-Matching[10]利用特权信息大幅提升了性能(目前最好的RefSR模型).明显地,教师网络的性能决定了学生网络的上限,但是现有利用特权信息的方法存在很大的问题.

根据局部相似性,教师网络提取HR高频特征因没有很好的约束导致学生网络对输入互补性不够准确,如图1所示.针对这个问题,本文提出了高效单图像超分辨率重建[11]:深监督对称蒸馏网络(DSSD).DSSD针对教师网络提取HR的高频特征不够准确提出了一个新颖的教师网络.为了约束中间层构造了高频特征递归模块(HFRM),这种递归方式对中间特征约束更强.为了使教师网络学习到足够准确的高频特征,提出了一种对称退化结构(SDM),SDM与HFRM一样进一步约束了教师网络中的SR.特别地,整个教师网络使用深监督加强对教师网络的约束,同时解决了教师网络难以训练的问题,使转移到学生网络的蒸馏知识更加准确.同时,本文为了更好适用于移动设备,选择只有13 k参数量和6 GFlops计算量的FSRCNN[5]网络.本文主要贡献:1)提出了一种新颖的通用蒸馏教师网络架构DSSD,该架构可以更准确地提取HR高频细节.2)其中构造了HFRM和SDM模块,目的是使教师网络更准确地提取HR高频特征.这两个模块采取递归方式并且结构简单.3)为了约束教师网络中提取的高频特征,采用深监督方法使教师网络蒸馏的知识可以与学生网络互补.

1 相关工作

1.1 单图像超分辨率

随着CNNs的发展,DONG等[3]首次提出基于CNNs的SISR方法SRCNN学习LR到HR的映射关系.之后的方法使用大量卷积层,利用密集和跳跃连接等保留更多的特征[12],与之前方法相比有了显著的提升.为了减少显存和运行时间,高效SISR被提出.对于显存高效方法,其本质是利用特殊设计的架构减少参数量.对于高效运行时间方法,其本质是较低的计算量,使用级联、多分支架构或组卷积等方式.但是这些特殊设计的高效SISR方法不能适用于移动设备.

1.2 知识蒸馏

知识蒸馏的提出是为了转移大模型(教师网络)软标签到小模型(学生网络)以增强小模型性能,它被广泛运用在网络压缩上[13].之后特征蒸馏被提出,可以从教师网络转移特征至学生网络,帮助学生网络更好地训练.利用特权信息的方法本质是使用额外信息,这在训练时需要额外成本但是在测试时不需要,被称之为广义蒸馏.广义蒸馏能够从教师网络转移到学生网络更多的信息以增强学生网络.特别地,变分信息蒸馏VID通过最大化教师和学生网络特征之间的互信息来转移知识,本文同PISR一样使用VID.

本文方法相关方法包括SRKD,FAKD和PISR.SRKD对教师网络和学生网络不同阶段进行蒸馏,考虑到教师网络与学生网络通道数不同,SRKD对中间特征的统计信息进行监督.FAKD改进了SRKD,在整个蒸馏的过程中FAKD与SRKD类似,区别在于FAKD在蒸馏损失方面提出了特征关联机制,取得了不错的效果.最近PISR被提出,PISR首次将HR作为特权信息引入SISR.PISR与SRKD、FAKD有本质的不同,其将HR直接输入教师网络,允许教师网络提取更多的信息转移到学生网络.

1.3 深监督

深监督是在训练神经网络时,针对中间隐藏层特征透明度不高以及深层网络难以训练的问题,对隐藏层进行再监督,促进网络更快更好收敛.DSN首次提出深监督概念,但DSN缺点在于使用的网络结构不够深,为了更好地训练深度网络,通过给神经网络的某些层添加了辅助监督来解决这个问题.此外,应用于分割的Unet++也适用深监督增强医学图像分割的性能.由于教师网络中间层较多,为了更好约束中间层,使中间特征能更准确地蒸馏到学生网络,本文采用深监督的方式,明显提升了DSSD的性能.

2 方 法

本文提出了一个新颖的深监督对称蒸馏教师网络.该网络利用特权信息训练教师网络,训练好的教师网络会帮助学生网络更好地重建.由于篇幅有限,本文以倍率4为例:使用X和Y表示LR和HR图像.其中′表示二分之一HR大小,″表示四分之一HR大小,ζ表示监督函数.

这种利用特权信息的对称结构类似于在课堂学习(训练阶段)时,不仅仅会有正确答案(HR),而且老师会解释其中的来龙去脉(HFRM),更进一步老师会告知错在哪里(SDM)以及该怎么修正(深监督).

2.1 教师网络

学生网络性能和教师网络性能直接相关,所以提高教师网络性能是关键.为此本文设计了一种对称教师网络结构如图2所示,该网络输入与PISR相同(HR).使用HFRM模块逐级下采样,每经过一个HFRM模块分辨率缩小二分之一,在递归2个HFRM模块后变成了HR分辨率的四分之一,与LR尺寸相同可以输入FSRCNN.在经过FSRCNN之后生成SR,将SR也采用与HFRM相同的下采样可以保证SR蕴含更多的高频特征,于是构造了SDM模块.特别地,SDM与HFRM模块在结构上相同,形成一种对称网络结构.同时为了转移更准确的高频特征,教师网络采取深监督的方式对中间层进行约束,取得了不错的效果.

2.1.1高频特征递归模块

本文构造了HFRM模块将HR递归投影到低维特征空间生成紧凑的高频特征.与机械式双三次下采样不同,基于CNNs的HFRM能学习到更多与LR互补的高频特征,然后从紧凑特征重构SR图像.递归HFRM受到深监督的约束可以提取更准确的高频特征.这一过程表示为式(1)、(2):

(1)

(2)

HFRM模块详细的结构如图3所示,k3n56s1为先经过一个3×3的卷积将3通道扩展到56(为了与FSRCNN保持一致),步长大小为1,激活函数为PReLU.第2个3×3的卷积步长为2进行下采样,之后再有两个3×3的卷积,先将56通道降至12通道,再变为3通道.经过2个HFRM生成与LR大小一致的紧凑特征交给FSRCNN进行重构,其中递归能很好地受到深监督约束.

2.1.2重建网络

本文目的是构造适用于移动设备的网络,既要很小的参数量也要很小的计算量,同时避免特殊设计的网络架构,所以FSRCNN是一个好的选择.教师网络中重建网络过程如式(3)所示:

(3)

2.1.3对称退化模块

为了保证经过FSRCNN重建的SR保留足够准确的高频特征,本文构造了SDM模块(图4),之前大部分方法主要是最小化SR与HR之间的损失而忽略了解空间的问题.对于LR来说,LR与SR是一种一对多的关系,利用SDM模块去约束其解空间的大小,保证SR具有更准确的高频特征.特别为了保持简单的结构,SDM模块与HFRM模块有着相同的构造,同样采用递归下采样的方式并受到深监督的约束,如式(4)、(5):

(4)

(5)

2.2 深监督

DSSD除了新颖的对称架构以外,还有一个明显的特点:它由多个相同的监督构成深监督约束.如图2所示,每个HFRM模块可以下采样2倍,于是4倍率需要递归2次.HR图像在经过HFRM模块之后会保留更多高频特征,这种特征会转移到学生网络.但是在保留的过程中,注意到学生网络的输入是双三次下采样得到的LR图像,尽管尽可能多地保留高频特征对学生网络有用,但是总体上不宜与LR偏差太大(局部相似性).为了解决这一个问题,本文使用深监督策略,每一级分辨率都受到监督,保证在保留高频特征的同时保证合理性.

(6)

(7)

(8)

(9)

(10)

(11)

(12)

(13)

2.3 学生网络

学生网络拥有与教师网络中重建网络(FSRCNN)相同的结构,但是输入不同.学生网络的输入为双三次下采样的LR,然后重建SR,如式(14):

(14)

初始化网络权重对学生网络的重建性能至关重要[14].使用教师网络的蒸馏知识初始化学生网络,可以将教师网络训练好的重建能力转移给学生网络,为后期优化提供了很好的起点.

尽管教师网络的蒸馏知识可以很好地初始化学生网络,但是由于学生网络的输入不同,因此需要使用重建损失和蒸馏损失进一步优化学生网络.其中式(16)中λS为10-6.

(15)

(16)

蒸馏损失将教师网络的蒸馏知识转移到学生网络,本文同PISR使用VID.目的是最大化教师网络和学生网络对应特征信息熵之间的互信息,

(17)

ζ(x)=ln (1+ex).

(18)

3 实 验

3.1 实验细节

本文使用DIV2K数据集训练,该数据集由800对LR和HR图像组成,其中LR是由HR双三次下采样得到.随机将HR图像剪切为192×192大小的块,同样的LR块取自对应比例因子的LR图像.同时使用数据增强技术,包括随机旋转和水平翻转.教师网络使用随机初始化训练,训练模型时,批处理大小为16,一共训练1 000轮,其中β1=0.900,β2=0.999.对于学习率大小,设置为10-3并且使用余弦退火技术将其减少到10-5.使用标准评估集,用峰值信号比(PSNR)和结构相似程度(SSIM)在Y通道上评估.本文运行于NVIDIA Titan RTX.

3.2 实验结果

附表Ⅰ比较了DSSD学生模型与最先进模型的性能,特别是高效SISR方法(*表示使用DIV2K重新训练,最好和最坏结果分别用下划实线和下划曲线标出).为了进行定量对比,评价指标为标准数据集在2、3、4×下的PSNR和SSIM,也包含模型的参数量计算量运行时间,其中计算量是在1 280×720大小的HR图像上测量.从附表Ⅰ可以看到:DSSD学生模型在所有比例因子上均优于PISR,甚至4倍率Set5超PISR接近0.12 dB,但是大数据集提升不如小数据集.因FSRCNN只有13 k参数量和6 GFlpos计算量,其重建能力在大数据集上已经捉襟见肘.同理,DSSD在较小的比例因子下的提升也有限,是因为较小的比例因子下丢失的信息较少,小参数量的模型重建能力有限.针对这个问题,本文后续在3.3.4讨论了较大模型的性能提升.

3.3 消融实验

消融实验对DSSD每个组件进行对比实验,以验证每个模块的作用,组合在一起之后会达到最好的效果.

3.3.1高频递归模块实验

为了验证HFRM的有效性,设计了表1对比实验:这里只探究HFRM模块,在设计此对比实验时去掉SDM模块以保证公平性.可以看到直接双三次下采样效果最差,因为机械式下采样就等于直接输入LR图像,导致教师网络和学生网络的输入相同.FSRCNN的PSNR为30.89 dB,直接双三次下采样方法训练本质上是用教师网络训练好的特征初始化学生网络.PISR的Encoder使用CNNs学到一部分LR没有的高频特征,但转移的高频特征不够准确(如图5对比所示).HFRM模块采用递归的方式约束中间特征,目的是可以学习到更准确更合理的高频特征.

表1 HFRM模块的结果比较

3.3.2对称退化模块实验

为了验证SDM的有效性,本文设计了表2对比实验.这里保留HFRM只改变SDM模块以获得公平的对比环境.从表2可以看出,去掉SDM就是表1最后一个实验.重点在表2的后3个消融实验:将SDM使用双三次下采样结果反而变差,原因是这种机械方式破坏了SR的退化结果,深监督会使SR向更坏的方向发展.而基于CNNs下采样限制SR的高频特征提高了SR的重建性能.最后是本文构造的SDM模块,在深监督的配合下进一步约束了SR高频信息的合理性,获得了更好的重建效果.

表2 SR的退化方式对比实验

3.3.3深监督实验

表3 监督对SR性能的影响

3.3.4教师网络通用性对比实验

为了探究DSSD教师网络是否具有通用性,本文挑选了几个参数量较大的轻量级SISR方法,见表4.图6为标准测试集在4视觉上的比较为665 k和591 k,*表示用DIV2K训练不设置蒸馏,可以看见DSSD在SISR方法都有提升.特别地,作用在较小的数据集Set5上提升较大,而在较大数据集上B100提升有限,这是因为FSRCNN参数量只有13 k已经到了性能极限.在增大参数量的网络VDSR上DSSD增强了更多,接近0.18 dB(Set5),而在IDN上的大数据集B100上增加了0.1 dB.这些实验验证了DSSD教师网络具有通用性.

表4 其他SR方法的定量结果

3.3.5视觉效果

图6展示了使用学生网络在标准测试集上的重建结果.可以清楚地看到学生模型提供了比原始基线模型(FSRCNN)和目前SOTA蒸馏方法PISR更好的性能.提出DSSD的目的在于探究如何更好地提取HR中包含的高频纹理细节,这些高频纹理细节可以使图像的边缘更清晰更明显,可以明显看见B100上数据集少了PISR的棋盘伪影,同时Urban100上数据集边缘明显变清晰锐利,Set14中barbara的重建效果明显纹理细节更多.这些效果图证明了DSSD确实具有很好提取HR高频细节的能力.

4 结 论

本文提出了一种新颖的高效单图像超分辨率重建方法:深监督对称蒸馏网络.针对教师网络提取HR高频信息不够准确这一问题,构造了高频特征递归模块和对称退化模块,有效提高了重建性能.在整个蒸馏的教师网络中,利用深监督能更好地约束中间特征使提取的特征更准确.消融实验很好地证明了本文方法的有效性,将在未来的工作中继续探索蒸馏方法中特征转移部分以进一步提高蒸馏方法的性能.

附 录

附表Ⅰ见电子版(DOI:10.16366/j.cnki.1000-2367.2023.06.007).