【2023 CCF 大数据与计算智能大赛】基于TPU平台实现超分辨率重建模型部署 基于QuickRNet的TPU超分模型部署

2023 CCF 大数据与计算智能大赛

《赛题名称》

基于QuickRNet的TPU超分模型部署

巴黎欧莱雅

林松

智能应用业务部算法工程师

中信科移动

中国-北京

[email protected]

团队简介

巴黎欧莱雅团队包含一个队长和零个队员。

队长林松,研究生学历,2019-2022在中国矿业大学(北京)攻读硕士学位,于2022年7月加入中信科移动公司,现在在智能应用业务部负责视觉AI算法的落地部署,是一名算法工程师,主要擅长视觉AI模型的训练、部署、优化,擅长x86、arm等平台的模型推理加速。

摘要

本文基于TPU超分模型部署提出了使用网络结构更简单、推理性能更高效、int8量化更友好的QucikSRNet来做为本次比赛的基础模型。

该模型已经被高通广泛应用到手机游戏、视频高清化、电视高清化、VR设备画面锐化等多个移动端领域,使得移动设备上实时超分辨率计算、热量和功耗都达到了不错的平衡。

在训练策略上,该模型使用了基于残差的超分辨率的训练初始化技巧,借鉴了类似重参数化的思想,在训练开始时将残差分支映射到权重为1的卷积嵌入到3x3的卷积,通过在随机初始化的3x3核权重上加一个中心值为1、其他值为0的权重核来模拟残差思想。

在推理策略上,本方案使用了膨胀推理的方案,并且使用了多模型多线程并行的生产者消费者模式,并且为了避免python语言层面的GIL问题,我将核心的前处理、模型推理、模型后处理统一封装到c++动态库中提供给python调用,推理速度上加速大于10倍。并且,通过本方案使用了内存池,显示的控制输入输出内存反复利用,避免了多次开辟和释放造成的耗时问题。

在平衡时延和性能上,本方案采取了网格搜索策略,对于模型的输出尺寸、模型的参数量大小、模型的量化方式进行了小范围内的搜索,最终选择输入尺寸192*192、模型参数1.7M、模型int8对称量化作为最终的方案,最终实现了模型平均niqe分数为5.0478和平均推理时间为38.3毫秒,获得初赛第三,复赛第二的成绩。

关键词

QucikSRNet,残差,量化,膨胀推理,生产者消费者,网格搜索

1 方案介绍

本方案使用了开源的QucikSRNet[1]作为基座模型,使用了其1.7M的大模型作为基准,并且分别针对其训练、推理进行了优化。初次之外本方案还对多种超参数进行了网格搜索,最终选择分数最高的组合作为最终的模型。

  1. 模型架构
    【2023 CCF 大数据与计算智能大赛】基于TPU平台实现超分辨率重建模型部署 基于QuickRNet的TPU超分模型部署_第1张图片

图 1:QuickSRNet模型结构图

目前移动端的SR算法结构都普遍比较简单,一般结构为全卷积+亚像素上采样(pixelshuffle、depth2space和transposeConv等),另外结合一个长连接将输入直接映射到输出,从而增加其训练稳定性,本方案的模型结构如图所示:

该模型参考了XLSR提出的clipped ReLU激活函数、参考了ABPN提出的基于锚点的长连接操作、参考了SESR和RepSR提出的重参数化方案、参考了ResNet提出的残差连接。下面是对参考方案的介绍:

1.1.1 Clipped ReLU

又叫做relu1,定义为:
在这里插入图片描述
(1)

相比relu,clipped_relu在两端都进行了截断,限制其特征值在0-1之间,这种操作可以有效降低模型在量化到int8时性能的下降。

1.1.2 基于锚点的长连接

将输入图像通过卷积改变通道后叠加到DepthToSpace模块前,相当于在输出之前叠加了一个插值操作,相比于通用型的SISR算法将输入上采样后叠加到最后,减少计算量的同时没有降低其性能。

1.1.3 重参数化

该模型吸收了残差连接和重参数化技巧的优点,对于每个卷积模块使用了权重初始化的新策略,由于普通的残差连接可以表示为:

在这里插入图片描述
(2)

上述操作可以重参数化为:

在这里插入图片描述
(3)

因此该模型将输入特征的映射路径作为权重为1的卷积嵌入到3x3的卷积中,具体实现可以通过在随机初始化的3x3核权重上加一个中心值为1、其他值为0的权重核。

除此之外,由于使用了BN保证训练过程的稳定,但是BN在训练和推理是不一致,因此在后期将BN参数冻结后继续训练,从而加速训练时的收敛。

  1. 推理策略

由于本次比赛提供的测试数据尺寸范围较大,在TPU模型的动态输入支持较差的情况下,本文测试了三种静态尺寸分块推理的方案,其中裁剪尺寸为192*192,具体方案分别如下:

1.2.1 无重叠区域的分块

对于任意尺寸的输入图像,当其尺寸小于192*192时,采用右下角补0的操作填充到192*192,当前尺寸大于192*192时,利用192*192的窗口滑动步距192,最终采样得到的n张图片依次前处理、模型推理、后处理最终粘贴到4倍超分的结果图中。

该方案没有多余的推理内容,推理速度较快,经过验证niqe值为5.1632,可视化图片发现分块边缘存在明显的痕迹如图:
【2023 CCF 大数据与计算智能大赛】基于TPU平台实现超分辨率重建模型部署 基于QuickRNet的TPU超分模型部署_第2张图片

图 2:无重叠区域分块可视化效果

1.2.2 有重叠区域的分块

相比于1.2.1的方案,有重叠区域的分块类似的但是采用滑窗步距为96,这种方案可以有效的减少1.2.1重叠区域带来的问题,但是如何融合重叠区域的结果是比较大的问题,采用均值的方案会让重叠区域采样的效果变得平滑,但是降低了锐化的效果,并且不同的滑窗步距会影响niqe值和推理的速度。

1.2.3 膨胀推理

为了结合1.2.1和1.2.2的优缺点,本方案提出使用膨胀推理多推少取的方案,将原始图像四周填充10个像素,采用192*192的滑窗,按照182的步长进行滑动,其中10为超参数。每次滑动得到的192*192区域经过模型推理后只取其中心的182*182的范围粘贴到输出图像中,为了避免填充0值带来影响,填充值采用了整个数据集的均值。下面是示意图:
【2023 CCF 大数据与计算智能大赛】基于TPU平台实现超分辨率重建模型部署 基于QuickRNet的TPU超分模型部署_第3张图片

图 3 膨胀推理示意图

使用该方法测试图片效果如图:
【2023 CCF 大数据与计算智能大赛】基于TPU平台实现超分辨率重建模型部署 基于QuickRNet的TPU超分模型部署_第4张图片

图 4膨胀推理可视化效果

  1. 生产者消费者模式

基于膨胀推理的方案会对输入图像进行分块,因此考虑使用生产者消费者模型将模型封装为一个模型池,每个生产者只需要负责自己线程的图像切块和前处理(归一化/通道转换),消费者只需要等待任意生产者完成前处理的信号然后马上启动推理,并将结果写入到输出图像中,整个流程如图所示:
【2023 CCF 大数据与计算智能大赛】基于TPU平台实现超分辨率重建模型部署 基于QuickRNet的TPU超分模型部署_第5张图片

图 5 生产者消费者模型

基于该模式本方案将前处理、模型推理、后处理分别给多个线程异步计算,测试服务器具有48核心可以大大加速整个推理的过程,相比于单线程的方案可以加速5倍以上。

  1. 网格搜索

由于本方案提出的膨胀推理、生产者消费者模型等存在如分块大小、线程数量、是否量化以及对称/非对称量化等超参数设置,本方案对于上述超参进行了简单划分,并且采用网格搜索策略搜索最优解,搜索部分结果表格如图:

输入尺寸 膨胀填充 量化方案 线程数量 最终分数
100*100 10 int8对称 16 3785
128*128 10 int8对称 16 3988
160*160 10 int8对称 16 4412
192*192 10 int8对称 16 5894
200*200 10 int8对称 16 4350
256*256 10 int8对称 16 3612
512*512 10 int8对称 16 3589
192*192 5 int8对称 16 5764
192*192 5 fp16 16 4755
192*192 10 int8对称 8 5656
192*192 10 fp16 8 5446
192*192 10 int8对称 32 5565
192*192 10 fp16 32 5311

表1:网格搜索部分结果展示

根据网格搜索结果可以看到,模型在小尺寸输入时会取得更好的结果,并且使用int8量化会带来更高的分数,在模型输入为192*192时,使用膨胀推理填充10个像素,使用16个线程可以达到最高的分数5894分。

  1. 其他优化

除了上述优化外,本方案还采用了线程池、内存池、算子融合、前后处理融合、模型预热、推理和计算指标分离等优化方法。

1.5.1 线程池

对于生产者消费者模型,本方案使用线程池进行了简单的实现,由于python在语言层面存在gil导致无法充分利用多核cpu,因此本方案基于c++的队列和信号量实现了上述功能,并且利用pybind将模型封装为动态库供python调用。

1.5.2 内存池

由于4倍超分输出的图像内存占用较大,为了避免反复开辟大的矩阵,以及避免在前后处理中临时矩阵所占用较大的内存,本方案设计了一个简单的内存池,在模型实例化时提前开辟较大的内存用于存储临时矩阵和输出矩阵,每次向python端返回内存的引用,这个操作可以加速5-10ms的时间。

1.5.3 算子融合/前后处理融合

本方案采取了常见的卷积+批归一化融合的操作,在将pytorch模型转换到onnx时完成算子融合。除此之外,本方案还将前处理中的归一化(除以255)操作融合到模型输入后的卷积权重中,方案时将这个卷积的权重除以255。为了简化图像BGR转RGB的操作,本方案将这个卷积的权重按照RGB通道进行了翻转,这样就避免了BGR图像需要再CPU中操作,模型可以直接输入BGR图像。

1.5.4 模型预热和推理计算指标分离

由于TPU在推理时前几次的速度较慢,因此本方案采取了预热30次推理空图来避免。

基线模型推理时会立刻计算niqe指标,由于推理时间相比于计算niqe指标较短,因此TPU会产生较长时间等待,为了避免这个问题,本方案采取先推理所有图片,存储推理耗时,然后在同一计算指标,这个方案使得TPU利用率可以达到80%以上,避免了等待问题。

致谢

感谢算丰官方免费提供的TPU服务器以及DataFountain提供的比赛机会,作为TPU模型部署的新手,这次比赛让我获得很多收获,也更加理解了超分任务在部署落地过程中的更多细节。

参考

[1] Guillaume Berger(2023年)。QuickSRNet: Plain Single-Image Super-Resolution Architecture for Faster Inference on Mobile Platforms。CVPR。

你可能感兴趣的:(大数据,超分辨率重建,人工智能)