AWS李沐：一种开源、端到端的新型人工智能框架编译器——NNVM编译器

谷歌AI是否能征服星际争霸？DT君在明年1月举办的全球新兴科技峰会，谷歌 DeepMind 研究科学家、《星际争霸》AI计划负责人 Oriol Vinyals 将发表主题演讲。预知详情，请点击图片！

文/李沐（亚马逊 AI 主任科学家）

原文：

https://amazonaws-china.com/cn/blogs/ai/introducing-nnvm-compiler-a-new-open-end-to-end-compiler-for-ai-frameworks/

你可以选择多种人工智能(AI)框架来开发人工智能算法，还可以选择多种硬件来训练和部署人工智能模型。框架和硬件的多样性对于维护人工智能生态系统的健康至关重要。然而，这种多样性也给人工智能开发人员带来了一些挑战。本文简要介绍了这些挑战，并介绍了一个可以解决这些问题的编译器解决方案。

让我们先来看看这些挑战，并向您介绍来自华盛顿大学和AWS研究团队，然后让您了解编译器是如何工作的。

三个挑战

首先，由于前端接口和后端实现之间的差异，从一个人工智能框架切换到另一个AI框架是非常重要的。此外，算法开发人员可能会使用多个框架作为开发和交付流程的一部分。在AWS上，我们有客户希望在MXNet上部署他们的Caffe模型，以享受Amazon EC2上的高速性能。Joaquin Candela最近在博客中写道，用户可能会使用PyTorch进行快速开发，然后部署在Caffe2上。然而，在将模型从一个框架转换到另一个框架后，有人抱怨二者的调试结果存在差异。

其次，框架开发人员需要维护多个后端，以保证智能手机芯片，GPU芯片等硬件的性能。以MXNet为例，它有一个可移植的，可以从头开始构建的C++实现。它还附带了目标相关的后端支持，比如cuDNN用于英伟达的 GPU芯片，MKLML用于英特尔的CPU芯片。保证这些不同的后端向用户提供一致的数字结果是很有挑战性的。

最后，芯片供应商需要为他们所构建的每一个新芯片支持多个人工智能框架。每个框架中的工作负载都以独特的方式表示和执行，因此，即使是一个简单的卷积（Convolution）操作，也可能需要以不同的方式定义。而支持多个框架需要花费工程师们大量的工作。

不同的人工智能框架和硬件给用户带来了巨大的好处，但是对于人工智能开发者来说，给最终用户带来一致的结果是非常困难的。幸运的是，我们并不是第一个面对这种问题的人。计算机科学在不同的硬件上运行不同的编程语言有着悠久的历史。解决这个问题的一个关键技术是编译器。受传统编译器技术的启发，Tianqi Chen，Thierry Moreau ，Haichen Shen，Luis Ceze，Carlos Guestrin以及Arvind Krishnamurthy这组来自华盛顿大学计算机科学工程学院的研究人员，联合了来自AWS AI团队的Ziheng Jiang，介绍了TVM堆栈技术来简化这个问题。

今天，AWS和来自华盛顿大学的研究团队共同发布了一种基于TVM堆栈技术的端到端编译器，叫做NNVM。这种编译器能直接将来自不同的深度学习前端的工作负载编译成优化过的二进制机器码。让我们来看一下传统的AI框架和NNVM编译器的结构对比：

我们观察到一个典型的人工智能框架可以大致划分为三个部分：

1、前端向用户展示了一个易于使用的界面。

2、来自前端的工作命令通常被表示成计算图（computation graph），它由数据变量和操作运算符组成。

3、不管是基本算术的操作运算符，还是神经网络层的操作运算符都被优化执行成对应不同硬件的操作运算符

而NNVM编译器基于TVM堆栈中的两个组件:处理计算图的NNVM组件和处理张量操作运算符的TVM组件。

• NNVM组件——计算图媒介表示堆栈 (intermediate representation (IR) stack)

NNVM组件的目标是将来自不同框架的工作命令表示为标准化的计算图，然后将这些高级计算图转换为执行图（execution graph）。将计算图以一种与框架无关的形式表达出来的想法，是受到来自Keras（一种深度学习框架）的层定义，以及来自Numpy（Python的一种工具）的张量操作运算符的启发。

NNVM组件还附带了一些称为Pass的例程，它们通过LLVM约定来操作这些计算图。这些例程可以在计算图中添加新的属性来执行计算图，或者修改计算图以提高效率。

• TVM组件——张量媒介表示堆栈

TVM起源于Halide语言，它的执行对象是计算图中的操作运算符，它把操作运算符优化成对应目标后端硬件的操作运算符。它与NNVM组件不同，它提供了一种与硬件无关的、对应特定领域的语言，以简化在张量索引级别中的操作符执行。TVM还提供了一些进程调度原语，例如多线程和缓存，从而优化计算以充分利用硬件资源。这些进程调度是硬件相关的，并且可以手动编码，也可以自动搜索优化过的模式。

NNVM编译器支持的前端框架和后端硬件

下图说明了NNVM编译器如何处理前后端之间的关系：

MXNet是通过将其计算图直接转换为NNVM组件中的计算图来支持的，Keras仍在发展中，但也以类似的方式支持。NNVM编译器还可以使用诸如CoreML之类的模型格式。

因此，任何能够使用这些格式的框架都可以使用这个编译堆栈。

TVM目前使用多个代码生成器来支持各种后端硬件。例如，它为X86和ARM等cpu架构生成LLVM IR。它也能够为各种GPU架构输出CUDA、OpenCL和Metal kernels。

NNVM编译器要添加新的支持很简单。对于一个新的前端框架，我们只需要将它的工作命令转换为定义了计算图和操作符规范的NNVM组件的格式。如果要添加新的后端硬件，我们可以重复使用TVM的操作符执行，然后只需指定如何有效地调度它们。

性能测试结果

我们用MXNet分别在GPU与CPU硬件后端上，以及ResNet18与MobileNet模型上，来与NNVM编译器进行性能对比测试。

这是在GPU上的性能结果测试对比图，可以看出NNVM编译器的耗时更少。

这是在CPU上的性能结果测试对比图，可以看出NNVM编译器的耗时更少。

结论

NNVM编译器能将一个高级计算图编译成优化过的二进制机器码，这个编译器基于TVM堆栈中的两个组件：NNVM组件提供了一个带有计算图优化例程的计算图和操作运算符的规范。TVM组件把操作运算符执行优化成对应目标硬件的操作运算符。而且，我们在两个完全不同的硬件上——ARM CPU和Nvidia gpu进行了性能测试对比，证明了这个编译器可以匹配甚至超越了目前最先进的计算性能。我们希望NNVM编译器能够极大地简化AI的前端框架和后端硬件的设计，并帮助用户在不同的前端和后端提供一致的结果。

-End-

编辑：Theo、黄珊

科技英语直播上线 | 跟着《麻省理工科技评论》每天学习科技英语！

点击左下方“阅读原文”进入小鹅通直播间

关注公众号：拾黑（shiheibook）了解更多

[广告]赞助链接：

四季很好，只要有你，文娱排行榜：https://www.yaopaiming.com/
让资讯触达的更精准有趣：https://www.0xu.cn/