深度学习模型inference优化之编译优化器

点击数：更新时间：2024-07-08 14:13:09

本文在 @蓝色的文章基础上补充修改
原文链接：

蓝色：手把手带你遨游TVM

人工智能无疑是目前风口浪尖的话题，大部分提到人工智能都关注在了算法，然而若要让人工智能落脚在实际应用中，工程化是至关重要的。本文脉络如下图所示，帮助大家理解。

模型的生产过程分为两步：Training部署、Inference部署。

对于Training部署，则是以业务数据为基础，设计相关的业务算法模型，最后使用一些流行的开源训练框架，如TensorFLow，PyTorch，Caffe等，或者自研的训练平台，如旷视的天元，清华的计图，训练得到一定精度的模型，市场上主流的训练框架基本上已经确定下来了。
对于Inference部署，则是将Training阶段生产的模型，部署到目标设备，嵌入到整个业务系统中去应用，但是对于部署平台则是“群雄逐鹿”。模型在Training得到后，可能会到多种多样的设备设备上运行，如下图所示，如Intel CPU / Intel GPU / ARM CPU / ARM GPU / NV GPU / FPGA / AI芯片等。

但是，想在这多种多样的设备中，都保持一个高效的Inference性能，其实是一件很有挑战的事情。

要达到模型能够在不同设备中都保持高效的Inference性能的目标是很不容易的。各种硬件设备的特性千差万别，如何保持一个统一的高效执行，是一个非常难做到的事情。而各大硬件厂商针对这样的情况都推出了自己的Inference 框架，相比TensorFlow等这样的框架孱弱的Inference性能，各大设备厂商的Inference框架性能都比较不错，如Intel的OpenVINO，ARM的ARM NN，NV的TensorRT等。但是这又带来更多的问题：

算子支持不够：训练框架也是多种多样的（如tensorflow，pytorch等），硬件厂商对不同训练框架支持力度是不同，会造相关算子没有及时实现；另外，不同硬件厂商的框架之间并不具备通用性，造成对不同训练框架的算子支持不全，尤其是像TensorFlow这种算子很多的。通常能在一个硬件厂商的Inference框架上跑的模型，不一定能在另外一个硬件厂商的Inference框架上跑。
业务环境多样性：对于业务开发来说也是非常痛苦的事情，在这个硬件上需要用这个Inference框架，在另外一个硬件上要用另外一个Inference框架，而且两者还没有统一的使用体验，算子支持也不一样，性能还不一定是最好的。

其实对于业务方来说，需要一个统一的Inference框架，然后在业务场景的各种硬件设备都能高效的运行，使用体验一致，只需要简单的根据硬件环境来配置相应的参数就好了。

各个硬件厂商针对自己硬件平台的Inference框架，使模型仅在自己的平台高效的运行。但是对于业务方（用户）来说，还是没很好的解决如何使自己的模型能在不同的硬件平台高效的运行的问题。

问题的解决可以类比编程语言。曾经出现了很多种编程语言，有很多种硬件，历史上最开始也是一种语言对应一种硬件，从而造成编译器的维护困难与爆炸。如下图所示。

而编译器解决了这个问题，其具体解决办法是这样的：抽象出编译器前端，编译器中端，编译器后端等概念，引入IR (Intermediate Representation)

编译器前端：接收C / C++ / Fortran等不同语言，进行代码生成，吐出IR
编译器中端：接收IR，进行不同编译器后端可以共享的优化，如常量替换，死代码消除，循环优化等，吐出优化后的IR
编译器后端：接收优化后的IR，进行不同硬件的平台相关优化与硬件指令生成，吐出目标文件

类似如下图所示的架构：

其实对于Inference框架也是类似的道理，只不过将编程语言，换成了不同的训练框架，相应的架构如下图所示：

把各种模型抽象看成各种编程语言，则需要引入一个新的编译器，如下图中红框（Compiler）所示，负责把这些不同框架的模型识别，然后吐出IR。

深度学习的模型构成是计算图，因此，可以将IR称为Graph IR。则Compiler的作用是输入不同训练框架训练的算子与模型，将其编译成Graph IR。

最后我们只需要实现对Graph IR的优化，即实现相应算子的优化，即可以将优化后的模型部署到不同的平台，最终结构如下图所示。

接下来就是实现的问题了：如何实现这样一个编译优化器，可以将不同训练框架训练得到的模型优化到不同硬件平台去高效运行？答案就是：NNVM compiler。如下图所示。

NNVM compiler基于此前发布的TVM堆栈中的两个组件：NNVM用于生成计算图，TVM用于映射张量运算。

NNVM：将不同框架的模型表示为标准化计算图，然后将这些高级图转换为执行图
TVM：提供一种独立于硬件的特定域语言，以简化张量索引层次中的运算符实现。另外，TVM还支持多线程、平铺、缓存等。

让我们给出一副更详尽，更漂亮的一张图：

NNVM compiler是亚马逊和华盛顿大学合作发布的开源端到端深度学习编译器，支持将包括mxnet，pytorch，caffe2，coreml等在内的深度学习模型编译部署到硬件上并提供多级别联合优化。速度更快，部署更加轻量级。支持包括树莓派，FPGA板卡，服务器和各种移动式设备和cuda，opencl，metal，javascript以及其它各种后端。

NNVM compiler可以将前端框架中的算子直接编译到硬件后端，能在高层图中间表示（IR）中表示和优化普通的深度学习算子，也能为不同的硬件后端转换计算图、最小化内存占用、优化数据分布、融合计算模式。

读完这篇文章，我相信，当大家处在将深度学习inference模型部署在不同硬件平台的泥沼当中时，现在已经有了自己的解决方案。