2023 CCF 大数据与计算智能大赛
《赛题名称》
基于QuickRNet的TPU超分模型部署
巴黎欧莱雅
林松
智能应用业务部算法工程师
中信科移动
中国-北京
团队简介
巴黎欧莱雅团队包含一个队长和零个队员。
队长林松,研究生学历,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: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参数冻结后继续训练,从而加速训练时的收敛。
- 推理策略
由于本次比赛提供的测试数据尺寸范围较大,在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膨胀推理可视化效果
- 生产者消费者模式
基于膨胀推理的方案会对输入图像进行分块,因此考虑使用生产者消费者模型将模型封装为一个模型池,每个生产者只需要负责自己线程的图像切块和前处理(归一化/通道转换),消费者只需要等待任意生产者完成前处理的信号然后马上启动推理,并将结果写入到输出图像中,整个流程如图所示:
图 5 生产者消费者模型
基于该模式本方案将前处理、模型推理、后处理分别给多个线程异步计算,测试服务器具有48核心可以大大加速整个推理的过程,相比于单线程的方案可以加速5倍以上。
- 网格搜索
由于本方案提出的膨胀推理、生产者消费者模型等存在如分块大小、线程数量、是否量化以及对称/非对称量化等超参数设置,本方案对于上述超参进行了简单划分,并且采用网格搜索策略搜索最优解,搜索部分结果表格如图:
输入尺寸 | 膨胀填充 | 量化方案 | 线程数量 | 最终分数 |
---|---|---|---|---|
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.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。
文章评论