【2023 CCF 大数据与计算智能大赛】基于TPU平台实现超分辨率重建模型部署 基于预训练ESPCN的轻量化图像超分辨率模型TPU部署方案
2023 CCF 大数据与计算智能大赛
《基于TPU平台实现超分辨率重建模型部署》
作品名：基于预训练ESPCN的轻量化图像超分辨率模型TPU部署方案
队伍名：Absofastlutely
蒋松儒
计算机科学与技术系 硕士
南京大学
中国-江苏
吕欢欢
计算机科学与技术系 博士
南京大学
中国-江苏
张凯铭
物理学系 本科
四川大学
中国-四川
团队简介
一个喜欢折腾神经网络和量子计算的无名小队，学艺不精脑洞不大，啥都只会一点，但没关系可以慢慢学的嘛……
摘要
本文介绍了一种在TPU平台上部署轻量化超分模型ESPCN的方法。介绍了三种图像分块策略、两种定义和搜索最优分块大小的方法；介绍了多线程加速、后处理滤波等模型性能优化技巧。最后在TPU平台上进行了模型部署和性能评估实验。
关键词
超分辨率，轻量化，TPU，模型部署
1 引言
2022年是图像生成式网络发展再度爆炸的一年。自Stable Diffusion开源，Embedding/Hypernetwork/LoRA微调、Depthmap/ControlNet条件化、AnimeDiff/Tune-A-Video连续帧化等等优化技术和下游应用也快速涌现。
但Stable Diffusion能直接产出相对优质图像尺寸是512~1024px，而我们时常需要更高分辨率如4k/8k的图像才能满足生产需求，因此超分辨率模型成了解决这个问题的一个可行桥梁方案。
总之就是，比赛要求参赛者在边缘设备TPU上移植并部署一个超分辨率模型。
2 解决方案
2.1 模型选择
考虑到TPU设备的计算性能相对较低，我们主要调研了轻量化的图像和视频超分辨率网络，并先后尝试了Real-ESRGAN[1]、NinaSR[4]、CARN、FSRCNN、ESPCN[2]等多种网络结构，最后选定了使用ESPCN。
2.1 图像分块
模型编译过程中需要确定输入张量的尺寸，虽然TPU文档里指出支持动态大小输入，但可用的SDK并未提供此功能，因此为了应对尺寸各异的输入图像，做统一尺寸的图像分块是必要的步骤。
2.1.1 分块策略
我们尝试了三种典型的分块策略：
- 朴素分块：即无重叠地裁剪后，各块分别推理，最后无重叠地粘接在一起。这种方式有最低的时间开销，但是会导致生成图像存在接缝。
- 重叠缝合：有重叠地裁剪后，各块分别推理，最后有重叠地粘接在一起，重叠区域按朴素平均、高斯羽化等方式做缝合。这种方式需要计算CNN模型的填充感受野，在重叠区域设置大于填充感受野时可以做到理论的无接缝，但缝合会产生额外的时间开销。
- 边缘裁剪：有重叠地裁剪后，各块分别推理并裁剪掉一圈边缘，最后无重叠地粘接在一起。这种方式也能实现理论的无接缝，额外计算面积比重叠缝合法小，但如果分片尺寸设置不合适，可能会导致更多的分片数量，反而增加额外计算量。
图1：三种分块策略（朴素分块、重叠缝合、边缘裁剪）
图2：朴素分块导致的接缝
考虑到时间效率优先，我们最终选用了朴素分块的策略。
2.1.2 最优分块大小
分块大小直接决定了分块数量，间接决定了模型处理处理的事件，因此合适的分块大小选取可以节省不必要的额外计算开销。我们探索了两种估算最优分块大小的方法：
- 最小化冗余计算量：考虑最小化分片导致的外向填充区域的面积。定义代价函数：
如图3所示，为代表原始图像大小的灰色区域，为所有红色分片总共覆盖的面积大小，为分片数量；是一个正则化常数，避免分片尺寸退化到无穷小，实验中取值为1。
图3：分片外向填充
使用双重退火、差分演化或者暴力穷举法最小化此代价函数都可以得到最优解 。
- 最小化单位计算费率：考虑最小化TPU处理张量中每个数据单元的时间开销。我们测试了ESPCN模型在TPU上处理不同尺寸输入的开销：
图4：TPU上处理不同尺寸输入的开销
观察可知，TPU设备上使用批处理并没有时间并行化效果，并且越小的张量尺寸计算费率越低。可以得到最优解 。
2.2 性能优化
为了进一步加速运行和提高图像质量，我们引入多线程处理和经典后处理滤波EDGE_ENHANCE和一个简化版本的UnsharpMask。后处理可以在CPU或者TPU上实现，视具体情况而定。
图5：引入多线程和后处理的稳定提升
3 部署和评估
模型部署：使用TPU-mlir[5]工具箱，将预训练好的ESPCN网络权重编译为TPU可加载的bmodel格式，量化类型设置为FP16，输入张量形状绑定为[1, 3, 128, 128]，部署运行时代码到TPU设备sophon BM1684[6]上。
使用下列计算公式来评估模型得分：
其中为平均自然图像质量评估度量[7]的得分，越小越好，而为模型的平均推理运行时间（单位为秒），同样越小越好。在两个各含有600张2k以下尺寸图像的数据集上测试表现如下：
| 模型 | 数据集 | runtime | niqe | score |
|---|---|---|---|---|
| Real-ESRGAN | test | 5.2462 | 3.9364 | 66.7269 |
| ESPCN | test | 0.1924 | 4.4465 | 1661.0888 |
| ESPCN | val | 0.0983 | 4.3761 | 3295.7186 |
对比可见，轻量化的ESPCN相比作为基准模型的Real-ESRGAN而言，虽然NIQE质量还很差，但凭借极高的推理速度，从而在综合得分方面有显著的提升。
致谢
感谢CCF BDCI平台提供的比赛机会和开放计算资源；感谢其他参赛队伍令人触目惊心的刷榜分数；感谢无数个太阳、朝露、地神的吐息。
参考
[1] Wang X, Xie L, Dong C, et al. Real-esrgan: Training real-world blind super-resolution with pure synthetic data[C]//Proceedings of the IEEE/CVF international conference on computer vision. 2021: 1905-1914.
[2] Shi W, Caballero J, Huszár F, et al. Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2016: 1874-1883.
[3] Lornatang. ESPCN-PyTorch [EB/OL] https://github.com/Lornatang/ESPCN-PyTorch
[4] Coloquinte. NinaSR: scalable neural network for Super-Resolution [EB/OL] https://github.com/Coloquinte/torchSR/blob/main/doc/NinaSR.md
[5] sophgo. TPU-MLIR [EB/OL] https://github.com/sophgo/tpu-mlir
[6] sophgo. BM1684X Introduction V1.7 [EB/OL] https://sophon-file.sophon.cn/sophon-prods3/drive/23/03/02/20/BM1684X%20Introduction%20V1.7.pdf
[7] Anish Mittal, Rajiv Soundararajan and Alan C. Bovik, Fellow, IEEE. Making a ‘Completely Blind’ Image Quality Analyzer. [EB/OL] http://live.ece.utexas.edu/research/quality/niqe_spl.pdf