结合编解码器与知识蒸馏的图像超分辨率方法

2022-12-28周兆京王晓茹姜竹青门爱东

重庆邮电大学学报(自然科学版) 2022年6期

周兆京，王晓茹，姜竹青，门爱东，马龙

(1.北京邮电大学人工智能学院，北京 100876；2.北京市经济管理学校信息技术系，北京 100089；3.中国人民解放军96962部队，北京 102206)

0 引言

近年来，图像超分辨率已成为一项重要的研究课题，它在目标检测、人脸识别和信息取证等方面有着重要的应用价值。超分辨率旨在学习一种非线性映射，从低分辨率图像中重建出高分辨率图像。传统插值算法主要是基于领域像素点进行加权计算来生成高分辨率图像，但仅能获得原图本身像素点的领域信息，无法生成新的高频信息，导致计算而得的高分辨率图像往往不够清晰。

随着深度学习的快速发展，研究者们提出了许多基于卷积神经网络的超分辨率方法[1-4]。SRCNN[2]是首个将卷积神经网络引入超分辨率领域的算法，取得了远胜传统算法的性能表现。SRCNN通过学习低分辨率图像和高分辨率图像的映射关系，可以重建出低分辨率图像中缺失的高频分量。增加卷积神经网络深度可以设计出卓有成效的超分辨率模型，进一步提高超分辨率网络的重建效果，但是其计算复杂度和内存占用量也急剧提升，直接在计算资源有限的设备端(如电视或手机)上实现它们是一个巨大的挑战。

为了解决这些难题，研究者们投入大量精力去研究如何压缩与加速神经网络模型[5]。大量的研究工作集中在网络剪枝、网络量化(将网络权值从浮点数量化为8比特整型值)、轻量化结构设计和知识蒸馏4个方面去优化神经网络的推理速度。其中，知识蒸馏是一种新颖的模型压缩方法，它通过将效果卓越的大网络中的知识传递到一个层数更浅、复杂度更低的小网络中去减轻计算负担，无需改变网络原本的结构特点。通过传递大网络的高级特征表达，小网络能够接收到比数据集里的标签更强有力的监督信息。

本文主要研究了知识蒸馏在超分辨率网络上的运用，通过大量实验论述了目前知识蒸馏方法在超分辨率上的局限性，并提出了一种新的知识蒸馏训练框架，基于编解码器的结构提取高分辨率图像中的先验知识，再将其传递给小网络，从而有效提升小网络的超分辨率重建效果，实现超分辨率网络的轻量化。

1 知识蒸馏方法

知识蒸馏[6]的主要思想是引导性能较弱、模型体积小的网络模仿性能更优、模型体积大的网络，以获得更优越的性能表现。一般而言，在知识蒸馏的训练模式中，网络层数深、复杂度高且性能更优的大模型被定义为大网络，网络层数浅、复杂度低且性能平平的小模型被定义为小网络。其关键在于如何定义大网络中的知识，以及如何传递知识。目前，知识蒸馏方法按照知识的定义可大致分为3类：基于softmax层输出的知识蒸馏[6]、基于特征的知识蒸馏[7-9]和基于特征关系的知识蒸馏[10-11]。这些方法在许多视觉任务上都取得了不错的实验结果。

然而，大多数知识蒸馏方法都是面向高级视觉任务而言的，如图像分类、目标检测等，而鲜有针对低级视觉任务的研究。为了实现超分辨率网络的压缩，本文将经典的知识蒸馏方法直接运用至超分辨率网络上，对比小网络采用蒸馏训练前后的重建效果，判断其方法的有效性，详细结果见表1。

表1 知识蒸馏在超分辨率网络的应用Tab.1 Application of knowledge distillation in super-resolution

表1采用文献[14]的EDSR作为基准网络，大网络采用与EDSR原论文相同的设置(包含32个残差模块、256个通道数)，小网络采用结构规模更小的EDSR(包含4个残差模块、64个通道数)，在Set5和Set14这两个标准数据集上进行了4倍超分辨率测试，使用峰值信噪比(point signal to noise ratio, PSNR)衡量重建图像的质量。PSNR是图像复原任务中使用最普遍的图像质量评估方法，其值越大，说明生成图像与真实图像越相近，图像质量越好。Soft-target[6]是基于softmax输出的知识蒸馏方法，其超分辨率任务输出的是高分辨率图像而并无softmax输出，实际实验中则修改为用大网络生成的图像监督小网络。AT[8]和FitNet[9]都是典型的基于特征的知识蒸馏方法，FSP[10]、CCKD[12]和SPKD[13]则是基于特征关系知识蒸馏的代表。

观察表1可得，这些知识蒸馏方法运用到超分辨率任务上，大多数都难以起到提升小网络性能的积极作用，只有基于输出图像的soft-target方法能带来一些改进。不同于分类等高级视觉任务，在超分辨率等低级视觉任务中，像素间的局部和全局关系尤为重要。考虑到这一特点，诸如FSP和AT等基于特征关系的方法是从特征图中抽象出更高级的特征表达知识，可能改变了图像本身的空间信息，从而影响了模型恢复图像的能力。而对于soft-target提供的加成，这可能是因为学习大网络的输出图像比直接学习真实图像更简单，一定程度上降低了模型的训练难度。图1展现了FSP与基准网络的训练过程对比情况，横轴是训练批量的迭代次数，竖轴是在Set5数据集上测试的PSNR结果，baseline是未经过知识蒸馏训练的小网络，FSP是采用FSP蒸馏训练的小网络。观察图1可得，知识蒸馏在训练前期能为小网络提供一个良好的助力，加快其收敛速度，但收敛后和基准网络逐渐趋于一致。

图1 FSP与基准网络的训练过程对比Fig.1 Comparison of the training process between FSP and baseline

综上可得，目前的知识蒸馏方法直接运用在超分辨率任务上效果甚微，仅能在训练前期起到加速效果。为了分析知识蒸馏在超分辨率等图像回归任务上效果不佳的原因，本文将超分辨率网络EDSR的中间层特征图进行了可视化，见图2。图2中，左图是输入图像，右图是EDSR首个残差块输出的特征图。观察EDSR的特征图可以发现，该特征图与原图轮廓保持着高度一致，这是由超分辨率任务本身特性决定的。分类网络的特征图往往趋向于关注目标的局部特征，其稀疏性较高，而超分辨率网络提取的特征仍然与输入图像趋于一致，每一部分细节都影响着成像质量，这无疑增大了通过蒸馏传递特征知识的难度。由此可知，要想知识蒸馏在超分辨率上发挥作用，关键在于大网络能提供更加有益于小网络训练的特征知识。

图2 超分辨率网络EDSR的特征图Fig.2 Feature map of EDSR

受文献[10]的启发，本文针对超分辨率任务提出了一种新的知识蒸馏训练框架，借助编解码器的结构，保留真实图像中的高频信息，为特征知识提供更精准、有用的特征，助力小网络的训练。

2 基于编解码器的特征知识蒸馏

2.1 构建大网络

本文的知识蒸馏训练框架如图3所示。图3中，HR表示训练集中的高分辨率图像，LR表示高分辨率图像对应的低分辨率图像。

图3 基于编解码器的特征知识蒸馏Fig.3 Feature knowledge distillation based on codec

不同于传统知识蒸馏中直接从已训练好的大网络中提取知识，本文知识蒸馏训练框架对大网络增加了一些结构上的训练约束，并使用高分辨图像作为大网络的输入对其进行训练。大网络分为编码器和解码器两部分，编码器对输入的高分辨率图像进行压缩编码，它将输入图像投影到一个低维特征空间中，生成更紧凑的特征，然后再输入解码器中重构出高分辨率图像，使大网络能为超分辨率任务提取更好的特征表示。

(1)

解码器的损失函数LD具体定义为

(2)

(2)式中，λ为退化损失的平衡系数；H、W分别表示高分辨率图像的高和宽。(2)式由2部分组成，第1部分是生成图像与真实图像的MAE损失，第2部分是生成图像退化后与解码器输入的MAE损失。

编码器、解码器共同进行训练，大网络能从高分辨率图像中提取高频信息，从而在网络中生成精确的特征知识，整个大网络的总的损失函数LT具体定义为

LT=βTLE+LD

(3)

(3)式中，βT是编码器损失函数的平衡系数。

2.2 构建小网络

小网络采用与大网络中解码器相同的结构，只是采用低分辨率图像作为输入，并使用解码器的网络权重对其进行初始化，为小网络提供一个良好的训练起点。虽然小网络和解码器具有相同的初始参数，但由于其输入不同，两者提取到的特征也大相径庭。

不同于解码器，小网络的损失函数不仅包含生成图像的重建损失以及退化损失，还包含蒸馏损失。这里的知识蒸馏则与传统知识蒸馏类似，是为了将解码器的特征知识迁移至小网络。

在蒸馏过程中计算中间层特征的分布，将特征图的分布信息定义为知识，通过最大均值差异(max mean discrepancy，MMD)[12]衡量大网络(解码器)和小网络之间的特征分布差异。以MMD作为蒸馏损失函数，引导小网络中间层的激活分布模拟大网络的激活分布。

(4)

(5)

(5)式中：‖·‖2是L2正则化；G是Gram矩阵，矩阵中每一项为gij=(fi)Tfj。Gram矩阵是特征图向量化后内积的结果，能反应特征之间的相关程度。

小网络总的损失函数LS可表示为

LS=Lsr+λLF+βSLdistill

(6)

(6)式中：Lsr为超分辨率任务常用的重建损失，由生成图像与真实高分辨率图像计算MAE损失而得；λ是退化损失函数的平衡系数，与大网络中的参数设置相同，λLF为退化损失，与解码器相同，将生成图像经退化后与输入图像计算MAE损失；βSLdistill为特征蒸馏损失，其中βS为蒸馏损失的平衡系数。

3 实验

3.1 数据集和评价指标

参照EDSR的训练设置[14-15]，本文使用DIV2K数据集进行训练，其中包含800张高分辨率图像，低分辨率图像通过对高分辨图像进行双三次下采样生成而得。每个训练批次大小为32张图片。

本文在Set5、Set14、B100和Urban100等标准数据集上评估提出的方法，并使用亮度通道上计算的PSNR作为评估指标。

3.2 实验分析

首先训练大网络，经过反复实验，设置超参数βT=10-4,βS=10-3,λ=0.1。表2所示为在不同βT值情况下训练小网络，并在Set5数据集上验证两倍超分的结果。在这部分实验中，退化支路没有引入其中，以免对编解码器的调参实验造成影响。当βT为0时，编解码器失去了压缩HR中高频特征的功能，成为了简单的线性映射。若参数βT设置太大，编码器损失函数将促使编码器生成的低维特征与低分辨率图像趋于同质。在这种情况下，本文提出的知识蒸馏框架将不能从高分辨率图像的先验知识中获益，并且小网络的特征蒸馏与传统蒸馏方法毫无差别，性能提升微乎其微。

表2 超参数βT的实验结果Tab.2 Results of balance parameters βT

经过大量实验得知，参数βT设置为10-4时，能在保证解码器学习超分辨率映射和编码器提取高分辨率的先验知识之间取得折衷，能通过特征蒸馏使小网络达到最佳效果。

小网络在不同数据集上进行2倍放大的蒸馏结果如表3所示，进行4倍放大结果如表4所示。经过蒸馏训练后，小网络在各个基准数据集上的性能都有所提升，在网络参数量和计算复杂度未增加的情况下，依靠本文提出的知识蒸馏方法，小网络PSNR在进行2倍超分辨率时能提升0.17～0.28 dB；在4倍超分辨率时能提升0.11～0.18 dB。与传统知识蒸馏方法相比，本文方法大大提升了小网络性能，能将知识蒸馏高效地运用在超分辨率任务上。这主要有以下两个原因：①大网络中编解码器结构能有效捕捉高分辨率图像中的高频信息，并通过特征蒸馏传递给小网络；②退化支路的约束缩小了超分辨率任务的解空间，加速了小网络的收敛。传统知识蒸馏方法未对大网络进行专门化训练，使得其蕴含的高频信息包含太多噪声，无法指导小网络训练。

表3—表4中的SRKD[18]和PISR[19]也是两种结合了知识蒸馏的超分辨率方法。其中，SRKD是基于传统知识蒸馏AT[9]的方式，通过学习大网络的特征图进行蒸馏，只是修改了AT中的知识定义，因此，它并不能脱离传统知识蒸馏在超分任务上的桎梏，PISR类似于本文方法，通过对小网络的预训练提取高分辨率图像中原有的先验信息来提高蒸馏效果。本文按照文献[18-19]的参数设置，在DIV2K数据集上完成了复现。结果表明，无论是在2倍还是在4倍超分辨率上，本文方法都取得了更高的PSNR指标。

表3 小网络在不同数据集上进行2倍放大的蒸馏结果Tab.3 Small network distillation results with 2x magnification on different data sets

表4 小网络在不同数据集上进行4倍放大的蒸馏结果Tab.4 Small network distillation results with 4x magnification on different data sets

将本文方法训练而得的小网络和文献[20-21]的超分辨率网络进行对比，从表3—表4可得，DRCN和MemNet的参数量都大于小网络，相应地计算复杂度更高，因而在各个数据集上它们都取得了比未蒸馏小网络更高的PSNR结果。但经过本文方法的训练，在不需对网络结构进行特殊设计情况下，蒸馏后的小网络并未增加运算复杂度或提高网络参数量，它的性能表现已经超过了DRCN，同时在耗时上也缩减至与MemNet相近的水准，这说明本文的知识蒸馏方法能针对超分辨率网络实现良好的轻量化效果。

3.3 消融实验

为了充分证明有效性，本文还对方法中的编解码器结构、退化支路、蒸馏损失函数进行了消融实验。在Set5上进行2倍放大，对PSNR进行定量计算分析，比较其超分辨率性能。

表5所示为在不同模块组合下训练而得小网络的消融实验结果。

表5 消融实验的结果Tab.5 Results of ablation studies

从表5可以看出，相比于传统蒸馏方法，编解码器结构使得特征蒸馏在超分辨率任务上更加有效，小网络受益于大网络解码器的网络权重，这为小网络提供了一个良好的训练起点，并且迁移了大网络的重构能力；MMD损失函数提供了比MAE更好的结果，基于MMD的蒸馏损失促使小网络和大网络的特征图保持一致分布，一定程度上避免了大网络特征图中的噪声影响；退化支路提供了更好的性能表现，能带来0.1 dB的PSNR提升，这表明缩小解空间有助于超分辨率任务。