Happy · 2023年02月23日

深入分析MobileAI图像超分最佳方案:ABPN

今天要介绍的MobileAI2021的图像超分竞赛的最佳方案,无论是PSNR指标还是推理速度均显著优于其他方案,推理速度达到了手机端实时(<40ms@1080P)

image.png

Abstract

尽管基于深度学习的图像超分取得前所未有的进展,但实际应用要求i越来越高的性能、效率,尤其是移动端推理效率。智能手机的升级迭代、5G的盛行,用户能感知到的图像/视频分辨率越来越高,从早期的480过度到720p,再到1080p,再到最近的1k、4k。高分辨率需要更高的计算量,占用更多的RAM,这就导致了端侧设备的部署问题。

本文旨在设计一种8-bit量化版高效网络并将其部署到移动端,整个设计过程如下:

  • 首先,我们通过将轻量型超分架构分解并分析每个节点的推理延迟,进而确定可利用的算子;
  • 然后,我们深入分析了何种类型的架构便于进行8-bit量化并提出了ABPN(Anchor-Based Plain Network);
  • 最后,我们采用量化感知训练(Quantization-Aware Training, QAT)策略进一步提升模型的性能。

我们所设计的模型能以2dB指标优于8-bit量化版FSRCNN,同时满足实际速度需求。

Method

接下来,我们从节点延迟测试开始,然后引出本文方案背后的思考,最后构建所提ABPN。

Meta-node Latency

由于我们的目标在于构建一种实时量化模型用于真实场景(比如实时视频超分)。我们需要做的第一件事就是构建可移植算子集并统计每个算子的耗时。

  • 我们将当前轻量型网络(如EDSR、CARN、IMDN、IDN、LatticeNet)进行分解构建初始算子集;
  • 我们在Synaptics Dolphin平台(专用NPU)上测试每个算子的延迟。

image.png

上述算子可以分为四大类:张量操作、卷积算子、激活算子、resize,见上表。从上表可以得出四个发现:

  • 近期的SOTA轻量架构使用的技术似乎难以在移动端部署
  • EDSR采用了大量的ResBlock,每个ResBlock会引入元素加,该操作甚至比高速优化的卷积还要慢;
  • CARN采用了全局与局部特征集成,每个集成过程包含一个concat与一个卷积,仅仅带来了0.09dB指标提升;
  • 由于大量的特征分离与拼接,IDN与IMDN同样存在端侧部署问题;
  • LatticeNet的部署问题更为严重,它采用了16个CA模块,每个CA模块包含一个元素加、一个元素乘、两个池化层,四个卷积,导致了过高的计算负担。
  • 另一个常见问题:它们都需要保存前面层的特征并采用控制数据流动。这种长距离依赖会导致RAM的低频处理,这是因为端侧内存非常有限。
  • 因此,我们将不考虑特征融合、特征蒸馏、组卷积以及注意力机制
  • 尽管卷积的参数量是卷积的9倍,但由于并行计算的缘故,两者的推理速度差别并不大。因此,我们采用卷积以得到更大感受野
  • 在激活函数方面,我们选择ReLU。这是因为它要比Leaky ReLu速度更快,而且i两者导致的性能差异非常小;
  • 由于HR与LR之间的坐标映射导致resize操作的推理速度过慢

Anchor-based Residual Learning

正如前一节所讨论的,能用的算子非常有限。为得到一个好的解决方案,我们深入分析了架构设计与INT8量化之间的相关性

据我们所知,其难度主要在于I2I(Image-to-Image, I2I)映射的高 动态范围,最直接的想法是生成低标准差权值与激活。有两种方式可以达成该目的:

  • 添加BN层:BN往往被集成在ResBlock中,尽管不会导致额外耗时与内存占用,但会导致0.2dB的性能下降。
  • 残差学习:近邻像素往往具有相似的值,很自然的一种选择就是学习残差。残差学习又可以分为以下两种:
  • ISRL:图像空间的残差学习
  • FSRL:特征空间的残差学习。

图像空间的残差学习在早期的工作(如VDSR, DRRN)中有得到应用,而特征空间的残差学习则更多在近期的SOTA方案(如SRGAN、IDN、IMDN)中得到应用并取得了稍优的性能。然而,我们认为:ISRL更适合于INT8量化

从前面Table1中可以看到:图像空间插值存在不可接受的推理耗时,甚至仅仅一次resize都无法满足实时需求。为解决该问题,我们提出了ABRL(Anchor-Based Residual Learning):它直接在LR空间复制每个像素9次为HR空间的每个像素生成锚点。受益于PixelShuffle层,所提ABRL可以通过一个concat+一个元素加操作实现。

image.png

上图给出了四种类型残差学习的区别所在,从推理耗时角度来看:

  • FSRL仅需要一个元素加操作,它的耗时为5.2ms;
  • ABRL包含一个通道拼接与一个元素加,总结耗时15.6ms,约为最近邻插值的四分之一。

所提ABRL有这样两个优点:

  • 相比FSRL,ABRL可以显著提升INT8量化模型的性能,提升高达0.6dB;
  • 多分枝架构可以通过并行加速,因此ABRL与FSRL的实际推理耗时相当。ABRL与FSRL的主要耗时源自RAM的访问速度慢

Network Architecture

image.png

image.png

Loss Function

在损失函数方面,我们采用了简单的L1损失,定义如下:
image.png

Experiments

在训练方面,图像块尺寸为,batch=16,优化器为Adam,初始学习率0.001,每200epoch减半,合计训练1000epoch。训练数据为DIV2K,在RGB空间评估性能。

QAT是一种流程的提升模型性能的量化技术且无额外推理耗时。我们设置初始学习率为0.0001,每50epoch减半,合计训练200epoch。QAT可以进一步提升0.06的B性能,此时INT8模型仅比FP32性能低0.07dB。

Residual Learning

image.png
上表对比了残差学习的性能、耗时。从中可以看到:

  • 对于FP32模型而言,FSRL模型取得了最佳性能,其他模型性能相当;
  • 对于INT8模型而言,不带残差的模型会出现严重性能下降(-1.93dB),FSRL模型会下降0.78dB,而ISRL则则仅仅下降0.13dB。因此,残差学习可以极大缓解INT8量化过程中的高动态范围问题,而ISRL变现优于FSRL

Test on Snapdragon 820

我们在Snapdragon 820的手机平台上,采用AIBenchmark软件测试了所提方案的CPU、GPU以及NNAPI耗时,结果见下表。

image.png

MAI2021 SISR Challenge

本文起初用于参加MAI2021图像超分竞赛,结果见下表。注:首次的提交的模型在模型尾部没有添加Clip操作,导致量化性能非常差(小于20dB);在竞赛结束后才解决了该问题并提交了校正后模型。受益于素体ABRL,所提方案取得了最佳PSNR指标,同时具有更快的推理速度。

image.png

来源:AIWalker
作者: Happy

推荐阅读

本文章著作权归作者所有,任何形式的转载都请注明出处。更多动态滤波,图像质量,超分辨相关请关注我的专栏AIWalker。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
6194
内容数
191
夯实深度学习知识基础, 涵盖动态滤波,超分辨,轻量级框架等
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息