NTIRE2023-RTSR-Track2 冠军方案详解

编辑 | Happy
首发 | AIWalker
链接 | https://mp.weixin.qq.com/s/JQ5g9yn_OdjR8hi_tWc4jA

arXiv:2305.02126 , cvpr2023 , code , video

本文了一个实时&轻量图像超分方案Bicubic++，它通过下采样模块降低图像分辨率以减少计算量，在网络尾部采用X6上采样进行图像重建，同时还构建了一个三阶段训练方案。在测试集上， 所提方案比Bicubic指标高~1dB，同时推理速度~1.17ms@RTX3090、2.9ms@RTX3070 (注：fp16精度，720p输入，4K输出)。此外，Bicubic++取得了NTRIE2023-RTSR-Track2 X3超分竞赛冠军，也是所有方案中最快的，所提方案与其他轻量化方案的性能对比可参考下图。值得称道的就是，Bicubic++具有与Bicubic相当快的处理速度，故很有机会成为新的工业标准。

评价标准

由于Bicubic++是面向NTIRE2023-RTSR竞赛而设计，在正式介绍方案之前，我们先介绍了一下该竞赛的评价标准。

其中，分别表示模型与Bicubic上采样的PSNR指标，t表示模型推理耗时(720p->4K)。此外，由于实时性要求，t必须小于30ms。计算公式就比较简单了，详细信息可参考calc_scoring.

def score(psnr, runtime, psnr_interp=31.697655):
    diff = max(psnr - psnr_interp)
    if diff == 0:
        return 0
    else:
        cte = 0.1
        return ((2 ** diff) * 2) / (cte * (runtime ** 0.5))

本文方案

上图给出了本文方案整体架构示意图，基于此，作者进行参数微调并尝试了不同的训练策略对最大化上述指标。

架构设计

在架构设计过程中，我们发现: 网络的速度并不直接依赖于参数量，而是与Activation相关。因此，更低分辨率的特征可以极大的降低Activation。此外，已有研究表明: 当与上采样操作协同学习时，下采样操作具有正向影响。因此，作者选择先采用stride卷积下采样，然后采用X6上采样以加速网络推理。在某种意义上讲，更低分辨率的特征提取可以认为是一种在空间上数据压缩方式，由于它与上采样协同训练，故该压缩是可逆的。

除了上面作者提到的stride卷积进行下采样外，作者还尝试了上图所提到的降采样方案，如Depth2Space、DWT。经综合考量后，最终选型stride卷积。

在网络结构方面，我们期望采用了最优的通道数以最大化模型推理速度。我们观测到：通道数与推理延迟并非严格的线性关系，见上图。由于可变通道数导致的延迟提升超过了其性能提升，我们决定网络通道数设置为相同通道数。结合上图同时保持算法实时性，作者将模型通道数设置为32。需要注意的是：为了得到最终的模型，作者以非最优通道数(34)模型开始，然后采用全局剪枝得到通道数为32的模型。

除此之外，作者注意到：卷积的bias项被忽视了，它会占据一定比例的整体延迟。因移除bias而导致的延迟减少要超越其导致的PSNR微弱降低。因此，作者在最终的模型中移除了bias参数。

上述发现可以概括为如下几点：

处理的数据量与推理延迟强相关，而非参数量；
模块间的通道数变动会降低延迟性能；
对于硬件来说，通道数与延迟之间存在一些 sweet spots，合理的利用可以一定程度提升模型性能；
bias会占据一定比例的延迟耗时。

Three-stage Training Pipeline

作者构建了一个三阶段训练范式:

Stage1 Slim: 采用更大的通道数进行训练。此时，ch=34, m=2, R=1, DS表示stride卷积。此时所有卷积均带有bias参数。

Stage2 Slimer: 全局结构化剪枝。基于上一步所得checkpoint参考上图执行全局结构化剪枝。需要注意的是： 此时剪枝准则并非参数大小或拟合程度，而是验证集上的PSNR指标。
Stage3 Slimmest: 移除卷积的bias参数项。在经过步骤二后，模型通道数从34减少到了32，加载上一步的checkpoint同时移除bias参数后对整个网络进行再次微调。