将OpenVINO预训练模型转为为ONNX,并使用TVM进行优化

OpenVINO是Intel推出的一款深度学习工具套件。OpenVINO带来大量的预训练模型,使用这些预训练模型可以快速的开发出自己的AI应用。

不过既然是Intel出的东西,自然少不了和Intel平台深度绑定。OpenVINO主要针对Intel的CPU进行优化。虽然也可以支持GPU,但支持的是Intel家的GPU。Intel家的GPU,应当不用报太多期待了。

为了支持更丰富的硬件类型,可以将OpenVINO自带的预训练模型 转为ONNX格式,然后在做其他处理。

OpenVINO模型导出为ONNX

OpenVINO优化后的预训练模型无法直接转换为ONNX。不过好在Intel有提供模型的训练和导出工具,利用OpenVINO的训练工具导出ONNX

OpenVINO用于训练和导出的库为: https://github.com/openvinotoolkit/training_extensions

具体的操作方式参见项目的具体说明文档。

对照人脸检测的文档,导出人脸检测对应ONNX模型: https://github.com/openvinotoolkit/training_extensions/tree/develop/models/object_detection/model_templates/face-detection

注:导出目录里有 export/export/alt_ssd_export/ 两种模型。其中 export/alt_ssd_export/ 包含了OpenVINO特有的实现,在转换为其他推理引擎模型时会失败,因此后续工作使用 export/ 中的模型。

使用TVM对ONNX模型进行优化

针对TVM的VM进行优化

对于存在动态shape的模型,TVM无法进行编译。很不幸的是OpenVINO中物体检测相关的模型都存在动态shape。在TVM无法编译的情况下,可使用TVM的VM进行执行。

  • 注:
    • 关于VM的相关内容请阅读: https://tvm.apache.org/docs/dev/virtual_machine.html
    • TVM的文档比较欠缺(特别是VM相关的内容)。不过好在项目还在快速迭代过程中,提交的issue很快就可以得到回复。
    • 根据测试,使用VM模式,在CPU上TVM的速度甚至比用 ONNXRuntime 还要慢不少。不知道是否是跑在虚拟机上的关系。

针对TVM进行编译和优化

如果你的模型可以正常编译,那就没必要采用VM模式了。直接编译理论上优化效果要好很多。这里采用的是TVM范例中给出的图片分类模型。

一个完整的模型优化和执行可以参考官方文档:Compiling and Optimizing a Model with the Python AutoScheduler

VM模式下加载和运行优化好的模型

加载前面导出的模型,并执行。

Leave a Reply

Your email address will not be published. Required fields are marked *