5.1.1. MLIR-迁移工具概述
TPU-MLIR是算能的智能视觉深度学习处理器编译器工程,该工程提供了一套完整的工具链, 其 可以将不同框架下预训练的神经网络, 转化为可以在算能智能视觉深度学习处理器上高效运算的模型文件 bmodel 。
代码已经开源到 github: https://github.com/sophgo/tpu-mlir 。
论文<https://arxiv.org/abs/2210.15016> 描述了TPU-MLIR的整体设计思路。
算丰系列智能视觉深度学习处理器平台仅支持BModel模型加速,用户需要通过使用算丰MLIR工具链,可以把其他框架下训练好的模型转换为BModel,使之能在算丰系列智能视觉深度学习处理器上运行。
MLIR已直接支持绝大部分开源框架(Pytorch、ONNX、TFLite、Caffe等)下的算子和模型,其他模型(TensorFlow、PaddlePaddle等)需要转换为ONNX模型后,再进行转换,更多的算子和模型也在持续开发中。
TPU-MLIR的整体架构如下图所示:
图 5.1 TPU-MLIR的整体架构
对于BM1688/CV186AH平台来说,TPU-MLIR支持float32、float16、int8、以及int4模型
另外,TPU-MLIR模型转换流程如下图:
图 5.2 TPU-MLIR模型转换流程
迁移模型需要在指定的docker执行, 主要分两步, 一是通过 model_transform.py 将原始模型转换成mlir文件, 二是通过 model_deploy.py 将mlir文件转换成BModel。
其中,生成FLOAT模型时,model_deploy.py 工具支持输出 F32/F16/BF16 三种浮点数据类型。
INT8/INT4模型, 则需要准备量化数据集、调用 run_calibration.py 生成校准表, 然后将校准表传给 model_deploy.py。如果INT8模型不满足精度需要, 可以调用 run_qtable.py 生成量化表, 用来决定哪些层采用浮点计算, 然后传给 model_deploy.py 生成混精度模型。模型移植样例可以参考5.1.3 INT8模型生成章节。
具体移植以及量化步骤可参考本章第二、三小节,更详细的问题可参考doc目录下的《TPU-MLIR快速入门指南》。