NTIRE2023-RTSR-Track2 冠军方案详解

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

NTIRE2023-RTSR-Track2 冠军方案详解_第1张图片

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相当快的处理速度,故很有机会成为新的工业标准。

NTIRE2023-RTSR-Track2 冠军方案详解_第2张图片

评价标准

由于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))

本文方案

NTIRE2023-RTSR-Track2 冠军方案详解_第3张图片

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

架构设计

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

NTIRE2023-RTSR-Track2 冠军方案详解_第4张图片

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

NTIRE2023-RTSR-Track2 冠军方案详解_第5张图片

在网络结构方面,我们期望采用了最优的通道数以最大化模型推理速度。我们观测到:通道数与推理延迟并非严格的线性关系,见上图。由于可变通道数导致的延迟提升超过了其性能提升,我们决定网络通道数设置为相同通道数。结合上图同时保持算法实时性,作者将模型通道数设置为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参数。
NTIRE2023-RTSR-Track2 冠军方案详解_第6张图片
  • Stage2 Slimer: 全局结构化剪枝。基于上一步所得checkpoint参考上图执行全局结构化剪枝。需要注意的是: 此时剪枝准则并非参数大小或拟合程度,而是验证集上的PSNR指标
  • Stage3 Slimmest: 移除卷积的bias参数项。在经过步骤二后,模型通道数从34减少到了32,加载上一步的checkpoint同时移除bias参数后对整个网络进行再次微调。

训练加速技巧

为尽可能充分利用硬件性能训练如此小的模型,所有的训练流程都需要进行适当调整。这是因为:模型前向与反向过程几乎不怎么耗时,验证集的访存反而会成为实际瓶颈。为加速验证,作者采用48张图进行验证,设置bs为8。考虑到访存瓶颈,作者在训练阶段将训练数据与验证数据均加载到RAM中已减少数据IO。

通过上述调整,1000epoch训练仅需1小时(单个Tesla V100)。此外,由于模型比较小,选择相当小的训练数据(如DIV2K)是可行的,也不会影响验证精度。

本文实验

NTIRE2023-RTSR-Track2 冠军方案详解_第7张图片

上表从不同维度对所提方案进行了消融分析,简单汇总如下:

  • 对比Q-T,可以看到:stride卷积具有更优指标;
  • 对比T&A,可以看到:相比ReLU,LeakyReLU可以带来明显PSNR提升且耗时几乎不变;
  • 对比A&E,可以看到:移除bias可以取得更高得分;
  • 对比1-4&E,可以看到:m=2,R=1可以取得最佳得分;
  • 对比PRB,可以看到:基于ch=34进行剪枝可以取得最高得分.

此外,为验证所提训练方案的有效性,作者还对比了A->G, E->H, A->F->I, A->E->H->J等不同优化路线。

NTIRE2023-RTSR-Track2 冠军方案详解_第8张图片

上表&图给出了所提方案同其他典型轻量化方案(如FSRCNN、ESPCN、ABPN、XLSR、SXCAT等)的客观与主观效果对比。

NTIRE2023-RTSR-Track2 冠军方案详解_第9张图片

上表&图给出了所提方案在NTIRE2023-RTSR竞赛数据集上的客观指标与主观指标。很明显,Bicubic++以显著优势超过了第二名方案。

推荐阅读

  1. ChatGPT聊图像超分,总结确实挺到位,哈哈
  2. 大核分解与注意力机制的巧妙结合,图像超分多尺度注意网络MAN
  3. CVPR2023 | 非局部注意力、局部自注意力、样条注意力协同助力图像复原达成新 SOTA!
  4. NAFNet :无需非线性激活,真“反直觉”!但复原性能也是真强!
  5. CVPR 2022 Oral | MLP进军底层视觉!谷歌提出MAXIM模型刷榜多个图像处理任务
  6. ELAN | 比SwinIR快4倍,图像超分中更高效Transformer应用探索
  7. CNN与Transformer相互促进,助力ACT进一步提升超分性能
  8. CVPR2022 | Restormer: 刷新多个low-level任务指标
  9. Transformer在图像复原领域的降维打击!ETH提出SwinIR:各项任务全面领先

本文由 mdnice 多平台发布

你可能感兴趣的:(后端)