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

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

《赛题名称》

基于QuickRNet的TPU超分模型部署

巴黎欧莱雅

林松

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

中信科移动

中国-北京

gpu@163.com

团队简介

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

队长林松，研究生学历，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. 模型架构
   

图 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，可视化图片发现分块边缘存在明显的痕迹如图：

图 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值带来影响，填充值采用了整个数据集的均值。下面是示意图：

图 3 膨胀推理示意图

使用该方法测试图片效果如图：

图 4膨胀推理可视化效果

1. 生产者消费者模式

基于膨胀推理的方案会对输入图像进行分块，因此考虑使用生产者消费者模型将模型封装为一个模型池，每个生产者只需要负责自己线程的图像切块和前处理（归一化/通道转换），消费者只需要等待任意生产者完成前处理的信号然后马上启动推理，并将结果写入到输出图像中，整个流程如图所示：

图 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。