旷视研究院 · 2022年11月08日 · 北京市

AI 模型编译器 MegCC 开源,让推理引擎体积变小小小小小……

目前社区中存在着不少个移动端深度学习推理框架(比如NCNN、MNN),它们为社区用户在移动端部署深度学习提供了相当多的便利,但是这些推理框架有一个共性问题:随着不断迭代以及性能优化,运行时库会逐渐增大,尤其是不同算子 fuse 时,会导致大量长尾算子,使 App 或 SDK 体积变得尾大不掉。

为了解决这个问题,由 MegEngine 团队开源的 MegCC 创新性地使用模型预编译方案,生成模型推理必要的代码,将与模型推理无关的代码去除,从而极大程度减小了推理引擎的体积。其主要方法是将传统框架运行时的必要步骤如计算图优化、Kernel 选择、内存分配等全部移到编译过程中,最大程度减少了 Runtime 时的二进制体积大小,并根据模型信息进行进一步的性能优化。

image.png

GitHub 开源地址:https://github.com/MegEngine/MegCC

方案特点

  • 伴随框架的迭代将不再增大推理引擎体积
  • 算子融合可以在编译时根据模型信息生成对应的代码
  • 模型编译时可以获得整个计算图的信息,以便继续进行极致的性能优化
  • 可以吸收社区在代码生成方面的经验用于为 MegCC 生成代码

不同于传统推理框架,MegCC 是一个真实的深度学习模型编译器,具备极轻量的 Runtime 二进制体积、高性能、方便移植、极低内存使用以及快启动等核心特点。用户可在 MLIR上进行计算图优化,内存规划,最后通过预先写好的代码模版进行代码生成。目前,MegCC 已支持 Arm64、Armv7、x86、risc-v 以及单片机平台。

image.png

MegCC 架构

使用方法及效果

使用 MegCC 完成模型部署只需要完成以下 3 步:

  1. 模型编译:编译 MegEngine 模型,生成运行这个模型对应的 Kernel 以及优化之后的模型。
  2. Runtime 编译:这个阶段会将 Runtime 和上一步中生成的 Kernel 一起编译成一个静态库。
  3. 集成到应用中:调用上一步编译的静态库的接口进行推理。

https://www.bilibili.com/video/BV1tg411B7dx/?spm_id_from=333.337.search-card.all.click&vd_source=3b44c0eca211dab0ebafeaf3160ba55e

详细操作指南见:https://github.com/MegEngine/MegCC/blob/main/doc/how-to-use-chinese.md

以 YOLOX 模型为例,运行效果如下图:
640.gif
从图中可见,MegCC 生成的推理程序在保证推理性能良好(模型测速结果为 670ms)的情况下,其大小可以达到 95KB。

未来计划

目前 MegCC 仅支持 MegEngine 模型作为输入,其他模型格式可以考虑转换到 ONNX,然后通过 mgeconvert 进行模型格式转换。

预计在未来 2 个月内,MegCC 将支持更多的模型格式编译。同时实现以下进阶功能:

  • 支持 ONNX 模型作为输入
  • 更多的 Kernel fusion
  • 支持更多的后端设备

大家在使用 MegCC 过程中有任何问题,欢迎随时提 issue 告诉我们,也欢迎提 PR 帮助 MegCC 变得更好。

来源:旷视研究院
作者:R

专栏文章推荐

欢迎关注旷视研究院极术社区专栏,定期更新最新旷视研究院成果。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
加入旷视:career@megvii.com
推荐阅读
关注数
7707
内容数
164
专注旷视研究院学术论文解读推送,涵盖计算机视觉,文字识别等
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息