DIPNet: Efficiency Distillation and Iterative Pruning for Image Super-Resolution

0. 简介

NTIRE 的全称为New Trends in Image Restoration and Enhancement Challenges，即“图像复原和增强挑战中的新趋势”，是CVPR(IEEE Conference on Computer Vision and Pattern Recognition)举办的极具影响力的计算机视觉底层任务比赛，主要涉及的研究方向有：图像超分辨率、图像去噪、去模糊、去摩尔纹、重建和去雾等。

其中在2023年，CVPR开展的NTIRE相关挑战有：

夜间摄影渲染（night photography rendering）；
从图像中估计深度（HR depth from images of specular and transparent surfaces）；
图像去噪（image denoising）；
视频着色（video colorization）；
阴影去除（shadow removal）；
视频增强的质量评估（quality assessment of video enhancement）；
双目超分辨率（stereo super-resolution）；
光场图像超分辨率（light field image super-resolution）；
图像4倍超分辨率（image super-resolution (×4)）；
360°全景图像和视频超分辨率（360° omnidirectional image and video super-resolution）；
镜头到镜头的散景效果转换（lens-to-lens bokeh effect transformation）；
实时4K超分辨率（real-time 4K super-resolution）；
高分辨率图像非均匀去雾（HR nonhomogenous dehazing）；
高效超分辨率（efficient super-resolution）。

同时，以上的这些挑战也蕴含着当前的一些研究难点及挑战，需要研究学者们集思广益，提出针对提升任务性能的想法，为共同解决近年来的难题贡献出一份力量。

本篇文章着重于NTIRE 2023高效超分辨率（efficient super-resolution）挑战赛的冠军方案的解读，学习能够提升任务的tricks，以期给相关的科研任务一些启发。NTIRE 2023高效超分辨率挑战赛的目标是以RFDN（AIM 2020 高效超分辨率冠军）为baseline，以尽量少的推理时间（runtime），参数量（parameters），计算量（FLOPs），激活值（activations）和显存占用（memory consumption）实现4倍超分，同时在DIV2K validation set上的PSNR至少要达到29.00dB。

比赛提供的数据集有DIV2K数据集和LSDIR数据集。DIV2K数据集包含1000张多样的2K分辨率的RGB图像，其中训练集有800张，验证集有100张，测试集有100张。LSDIR数据集有86991张高分辨率高质量的图像，其中训练集有84991张，验证集有1000张，测试集有1000张。

1. 摘要

最近对高效超分辨率的研究主要集中在通过各种网络设计来减少参数量（the number of parameters）和浮点数运算量（floating-point operations）。尽管这些方法能够减少参数量和计算量，但是并不一定能够真正减少推理时间（running time）。

针对这个问题，作者提出了一个多阶段轻量级网络增强方法，能够使得轻量级网络获得优异的性能。特别地，作者使用经过增强的高分辨率结果（实际上是其他网络的1倍超分结果）作为训练时的额外监督，从而提升轻量级学生网络的学习能力。

学生网络训练收敛之后，使用重参数化技术（reparameterization）技术和迭代式网络修剪（iterative network pruning） 来简化网络的结构。

同时，作者还采用了结合多锚点蒸馏（multi-anchor distillation）和渐进式学习（progressive learning） 作为高效轻量化网络训练策略，使得轻量化网络获得杰出的性能。

最后，该方法在保持极具竞争力的超分结果的同时，在NTIRE 2023高效超分辨率挑战的所有方案中，推理速度是最快的。结果表明，该方法在极具代表性的DIV2K数据集上（在客观和主观方面）取得了相当好的性能，同时推理速度更快，网络参数量更少。

2. 引言

最近的单幅图像超分（SISR）方法（如RFDN）在重建图像方面取得了巨大的成功。然而，这些方法通常需要大量的计算资源，使得很难部署在真实世界应用的资源受限的设备中。

为了解决这个问题，越来越多研究者开发出具有高效推理速度的SISR模型，同时要在重建图像质量与计算量上取得更好的权衡（trade-offs）。之前的研究尝试通过减少模型参数或浮点操作来提高效率，具有权重共享的递归网络通常用来减少参数量，但是由于其复杂的图拓扑结构，并不一定能减少计算量和推理时间。同样的，通常也会使用深度可分离卷积（depth-wise convolutions）、特征划分（feature splitting）和混洗（shuffling） 来减少FLOPs，但这也不一定能提升计算效率。

因此，作者考虑从其他的角度解决问题，即无需额外的网络设计，只通过更好的训练策略，就能达到高效超分。如下图所示，作者提出的策略使得超分模型更快更小。直接训练好一个小网络是困难的，但是我们可以通过先训练一个大的教师网络，再让教师网络通过知识蒸馏（knowledge distillation）的方式指导学生网络的学习。但是由于教师网络和学生网络的学习能力存在巨大的差距，所以学生网络很难通过直接蒸馏的方式学到足够的高频信息。受到HGGT的启发，作者使用的高分辨率（HR）图像不是原始数据集中的HR图像，而是经过增强的HR图像。增强后的图像能够提供更丰富的高频信息，能够帮助学习能力有限的学生网络更轻松的学习。受到RepVGG的启发，在设计学生网络的过程中，作者设计了顺序分支，并行分支和残差连接。当学生网络训练收敛后，使用重参数化的技术简化成一个轻量级结构。这一操作增强了模型的学习能力，但又没有引入额外的模型复杂度。为了充分使用教师网络的指导能力，作者使用多层级蒸馏策略，即在网络的不同结点处设置锚点，使用锚点处不同层级的特征来进行蒸馏。

通常，为了训练的方便，我们使用相对较小的图块（patch）进行训练。然而，近期研究表明这样的方式将会造成训练和测试时的输入分布不一致，使得网络在测试时不能高效地进行一些全局操作。因此，为了使得模型在测试阶段表现得更好，训练阶段和测试阶段的输入尺寸应该尽可能接近。然而，在训练时直接使用大尺寸的图块会非常耗时，不能使用更大的图块也会进一步影响训练的稳定性。同时，在表达能力有限的网络中，使用大的图块，并不利于挖掘全局的信息。因此作者采用了渐进式学习的方法，在训练过程中逐渐增大输入的图块尺寸，能够取得很好的结果。训练好的学生网络仍有一些不重要的冗余参数，所以作者迭代式地修剪模型，进一步压缩模型尺寸。

作者的贡献总结如下：

首先，提出了使用增强的HR图像来增强轻量级网络的学习能力；
作者提出了一个结合蒸馏，渐进式学习和修剪的多阶段轻量化训练策略；
大量实验表明论文所提方法的有效性，并且该方法在推理时间和模型大小上超过了NITRE 2023高效超分挑战赛的所有其他对手。

3. 相关工作

相关工作中回顾了一些主流的单幅图像超分和高效图像超分领域的方法。

3.1 单幅图像超分

Dong等人提出了第一个基于CNN的超分方法——SRCNN。
Kim等人提出了一个更深（20层）的网络来提高超分性能——VDSR。
Lim等人使用局部和残差连接，提出了EDSR。
RCAN引入了通道注意力和二阶通道注意力，可以利用特征的相关性来提高性能。
一些工作提出了改进的损失函数来学习真实的图像细节——VGG loss、perceptual loss和GAN loss。
基于Transformer的超分方法越来越流行，能够获得更好的性能——SwinIR和HAT。
然而，大部分方法需要大量的计算资源，参数量、FLOPs和推理时间都比较大，不能在边缘设备上进行部署和应用。

3.1 高效图像超分

高效图像超分旨在减少计算量和参数量，同时取得更快的推理时间和保持良好的性能。在真实世界的SR模型部署中，边缘设备的算力通常十分有限，因此也有大量的研究工作集中于高效的超分技术。

NTIRE和AIM几乎每年都会开展高效图像超分赛道，促进这方面的相关研究（大家可以好好留意一下这两个比赛）。
CARN（2018）使用了局部和全局级联机制，从而获得轻量级的SR网络。
IMDN（2019）通过构造级联信息多蒸馏块，设计了一个信息多蒸馏网络来提取层级特征。
RFDN（2020）引入了特征信息块来进一步提升网络性能，使用 $1\times1$ 卷积层实现维度的转换。
RLFN（2022）在RFDN的基础上，探究了其速度瓶颈，通过移除层级特征连接来提高推理速度。此外，RLFN提出了特征提取器，能够提取边缘和纹理的更多信息。RLFN是NTIRE 2022高效超分挑战赛的冠军。

4. 方法

作者提出的高效蒸馏和迭代式修剪SR网络，称为DIPNet，该网络包含四个主要模块，分别为：

重参数化残差特征块（reparameterization residual feature block, RRFB）；
基于模型引导的GT增强（Model Guided Ground-truth Enhancement）；
多锚点特征蒸馏（Multi-anchor Feature Distillation）；
迭代式修剪策略（Iterative Pruning Strategy）。

4.1 重参数化残差特征块

遵循RFDN和RLFN，作者同样使用信息蒸馏网络来重构高质量的SR图像。基于RLFN的RLFB块，作者在块中引入了可重参数化的拓扑结构。下图中，（a）是原始的RLFB块，在训练过程中，作者将其拓展成重参数化残差特征块（RRFB）的结构，如（b）所示。RRB的结构如（c）所示，在保证计算效率的同时，挖掘了优化过程中复杂结构的潜力，因为在推理过程中，RRB与单个 $3\times3$ 的卷积在计算上是等价的。

4.2 基于模型引导的GT增强

现在大多数方法在训练阶段都直接使用原始的HR图像。然而，根据HGGT提到的那样，原始HR图像的感知质量可能还是不够好，没有足够好到能够很好地监督训练。受到HGGT的启发，作者提出了一个模型引导GT增强策略，用来增强HR图像的（感知）质量。首先，作者训练了一个大型的混合注意力Transfomer（即HAT），用于实现1倍超分，用来产生增强的HR输出，记作 $I_{enh}$ 。然后使用低分辨率图像 $I_{LR}$ 和增强的HR图像 $I_{enh}$ 监督4倍超分的训练。不同于HGGT，作者没有进行人工的图块选择，保留区域较为平坦的图块，因为手工的选择通常能使得产生的结果视觉效果更好，但是对于客观指标却没有提升。下图展示了某些图块经过增强后的质量会优于原始的GT。

4.3 多锚点特征蒸馏

为了进一步提升轻量级模型的性能，作者提出了多锚点特征蒸馏方法。如下图所示，多锚点特征蒸馏包含两个阶段，第一个阶段是训练一个大的教师模型HAT，该模型记作 $\tau$ 。值得注意的是，在教师模型4倍超分训练过程中，作者使用增强的HR图像 $I_{enh}$ 用于最小化损失，即最小化 $||\tau(I_{LR})-I_{enh}||_1$ 。

在训练完教师网络之后，作者对学生网络进行多层级的蒸馏，把学生模型记作 $S$ 。一旦学生模型训练收敛，就使用对RLFB结构使用重参数化技术。在蒸馏过程中，作者使用教师模型 $\tau$ 的四个不同层级提取出的特征图，来监督学生模型 $S$ 的四个块的学习。

同样地，作者也是用教师模型地输出的伪GT和增强后的GT进一步监督学生模型 $S$ 的学习。在完成蒸馏后，作者使用渐进式学习的策略微调学生模型。也就是在训练过程中逐渐增加输入图块的尺寸，同时使用 $L_2$ 损失来监督训练，直到学生模型完全收敛。

4.4 迭代式修剪策略

最后，作者迭代式地修剪重参数化之后的学生网络。受到AGP的启发，在模型训练过程中，作者使用 $L_2$ 滤波修剪的方式对网络进行迭代式地修剪。直到网络不能高效地进行预测后，作者才停止对网络的修剪，并把上一次修剪的网络作为最后的网络。

5. 实验

5.1 实验设置和细节

训练集：DIV2K。

测试集：Set5、Set14、BSD100、Urban100和Manga109。另外还有NTIRE 2023高效超分挑战赛提供的DIV2K和LDSIR的测试集。

评估指标：RGB通道的PSNR和SSIM。另外还有参数量（the number of parameters）、计算量（GFLOPs）和推理时间（inference time）。

比较模型：CARN、IMDN、RFDN、RLFN等等。

训练设备：NVIDIA 2080ti。

数据增强策略：随机翻转和旋转。

优化器：Adam。

训练策略：在训练教师网络和蒸馏学生网络的时候，初始学习率设置为 $10^{-4}$ ，每100k个迭代次数学习率减半，使用 $L_1$ 损失进行监督。当使用渐进式学习策略微调学生网络时，初始学习率设置为 $2\times10^{-5}$ ，每200k个迭代次数学习率减半，输入的图块大小逐渐增大，依次为 $[64, 128, 256, 384]$ 。在迭代式修剪的过程中，修剪率为0.05，一共重复3次。在每轮修剪之后， $I_{LR}$ 和 $I_{enh}$ 都用于微调过程的监督， $384\times384$ 的图块作为微调过程中的输入。

5.2 模型复杂度

下表展示了在维持29.00dB的PSNR时，作者提出的DIPNe使用的资源最少，推理时间快过RFDN，迭代式修剪策略使得网络在最少的开销中，能够获得更高的速度增益，远远快于RFDN。

5.3 主观比较

5.4 客观比较

5.5 消融实验

Original GT vs. Enhanced GT

多阶段特征蒸馏

迭代式修剪策略

5.6 NTIRE 2023高效超分挑战赛

6. 结论

在本文中，作者提出了一个新颖的高效单幅图像超分的方法，该方法不是只依赖于网络结构的设计，而是通过改进训练策略来实现的。特别地，作者使用增强后的GT图像作为训练时额外的监督，结合蒸馏，渐进式学习和修剪提出了多阶段轻量级训练策略。实验表明了所提方法的有效性，针对推理时间和模型大小，在NTIRE 2023高效超分挑战赛中取得了SOTA的性能。论文的贡献在于：使用增强后的GT图像提升轻量级网络的学习能力，提出了多阶段轻量化训练策略。

最后感谢小伙伴们的学习噢~