paper: https://arxiv.org/abs/2305.02126
code: https://github.com/aselsan-research-imaging-team/bicubic-plusplus
本文了一个实时&轻量图像超分方案Bicubic++
,它通过下采样模块降低图像分辨率以减少计算量,在网络尾部采用X6上采样进行图像重建,同时还构建了一个三阶段训练方案。在测试集上, 所提方案比Bicubic指标高~1dB,同时推理速度~1.17ms@RTX3090、2.9ms@RTX3070 (注:fp16精度,720p输入,4K输出)。此外,Bicubic++取得了NTRIE2023-RTSR-Track2 X3超分竞赛冠军,也是所有方案中最快的,所提方案与其他轻量化方案的性能对比可参考下图。值得称道的就是,Bicubic++具有与Bicubic相当快的处理速度,故很有机会成为新的工业标准。
评价标准
由于Bicubic++是面向NTIRE2023-RTSR竞赛而设计,在正式介绍方案之前,我们先介绍了一下该竞赛的评价标准。
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参数后对整个网络进行再次微调。
训练加速技巧
为尽可能充分利用硬件性能训练如此小的模型,所有的训练流程都需要进行适当调整。这是因为:模型前向与反向过程几乎不怎么耗时,验证集的访存反而会成为实际瓶颈。为加速验证,作者采用48张图进行验证,设置bs为8。考虑到访存瓶颈,作者在训练阶段将训练数据与验证数据均加载到RAM中已减少数据IO。
通过上述调整,1000epoch训练仅需1小时(单个Tesla V100)。此外,由于模型比较小,选择相当小的训练数据(如DIV2K)是可行的,也不会影响验证精度。
本文实验
上表从不同维度对所提方案进行了消融分析,简单汇总如下:
- 对比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竞赛数据集上的客观指标与主观指标。很明显,Bicubic++以显著优势超过了第二名方案。
来源:AIWalker
作者: Happy
推荐阅读
- AsConvSR | NTIRE2023-RTSR-Track1冠军方案
- YOLO-NAS | YOLO新高度,引入NAS,出于YOLOv8而优于YOLOv8
- ChatGPT是如何看待YOLO系列算法的贡献呢?
- APE | 用于可扩展的单图像超分中自适应“块早退”
- 比MobileOne还秀,Apple将重参数与ViT相结合提出FastViT
本文章著作权归作者所有,任何形式的转载都请注明出处。更多动态滤波,图像质量,超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。