陈天奇团队发布TVM：把深度学习部署到手机、树莓派等更多硬件

允中 编译整理
量子位 出品 | 公众号 QbitAI

△ 陈天奇，华盛顿大学计算机系博士生，此前毕业于上海交通大学ACM班。XGBoost、cxxnet等著名机器学习工具的作者，MXNet的主要贡献者之一。

DMLC项目发起人陈天奇今天早间宣布推出TVM。

所谓TVM，按照正式说法：就是一种将深度学习工作负载部署到硬件的端到端IR（中间表示）堆栈。换一种说法，可以表述为一种把深度学习模型分发到各种硬件设备上的、端到端的解决方案。

陈天奇在微博上表示，TVM和之前发布的模块化深度学习系统NNVM一起，“组成深度学习到各种硬件的完整优化工具链”。

同在DMLC小组的刘洪亮（phunter_lau）进一步在微博上解释了这个工作的意义：“TVM可以把模型部署到不同硬件，比如群众常问的能不能用AMD的GPU，用FPGA怎么搞，TVM提供这个中间层有效解决这个问题”。

随后陈天奇也补充说：

除了比较明显的多硬件支持，更重要的是支持比较方便的自动调优和轻量级部署。比如我们有同学可以在一些workload可以达到和cudnn差不多的效果，而且同样的东西可以迁移到其它非cuda设备。

非常建议大家尝试一下。

而在reddit上，刘洪亮形象的比喻称：以后可以让树莓派来找猫~

关于TVM的官方介绍，量子位尝试把主要内容编译如下。查看原文可以点击页面左下角“阅读原文”按钮。

作者：Tianqi Chen(project lead), Thierry Moreau(hardware stack), Ziheng Jiang(graph compilation), Haichen Shen(gpu optimization)

深度学习已经变得无处不在、不可或缺。

此次变革的推手之一，是可扩展的深度学习系统，例如TensorFlow、MXNet、Caffe以及PyTorch等。大多数现有系统只对部分服务器级GPU进行了优化，如果想部署到手机、物联网设备以及专用加速器（FPGA、ASIC）等平台，还有大量的工作要做。

随着深度学习框架和硬件后端数量的增加，我们提出一个统一的中间表示（IR）堆栈，用来弥合深度学习框架和硬件后端之间的距离。

我们很高兴的宣布推出TVM来解决上述问题。TVM是一个全新的框架，可以：

• 为CPU、GPU和其他专用硬件，表示和优化常见的深度学习计算工作负载

• 自动转换计算图以最小化内存占用，优化数据布局和融合计算模式

• 提供端到端编译，从现有的前端框架到裸机硬件，直到浏览器可执行的javascript

在TVM的帮助下，可以轻松在手机、嵌入式设备甚至浏览器上运行深度学习的工作负载，而不需要额外的工作。TVM还为许多硬件平台上的深度学习工作负载，提供统一的优化框架，包括依赖于新计算基元的专用加速器。

我们采用了编译器界的共同理念，提供两个中间表示层，以有效地将高级深度学习算法降低到多种硬件后端。

在这次放出的版本中，开源的TVM软件包提供x86、ARM、OpenCL、Metal、CUDA和JavaScript的优化基元。我们正积极的致力于增加对专业硬件加速和Nvidia GEMM优化的Volta架构的支持。

技术细节

TVM堆栈的目标，是提供一个可重复使用的工具链，来将高级神经网络描述从深度学习框架前端，向下编译为多个硬件后端的低级机器代码。

以Apache MXNet作为前端案例，下面的代码演示了如何使用TVM将深度学习模型的高级描述编译为针对目标硬件定制的优化可执行模块。

这件事的挑战在于支持多个硬件后端，同时将计算、内存和能量足迹保持在最低水平。我们借鉴了编译器界的智慧，构建了两级中间层：其中一层是NNVM（用于任务调度和内存管理的高级中间表示），另一层是TVM（用于优化计算内核的富有表现力的低级中间表示）

堆栈的第一级是基于计算图的表示。计算图是一个有向无环图，用节点表示计算，用箭头表示数据流关系。大多数现有深度学习框架都采用这种方法，包括TVM堆栈中的NNVM图表示，TensorFlow XLA以及英特尔的Ngraph。

图优化框架可以支持很多强大的优化。例如，我们提供了一个次线性内存优化功能，允许用户在单个GPU上训练1000层的ImageNet ResNet。

然而，我们发现仅基于IR的计算图不足以解决支持不同硬件后端的挑战。因为单独一个图形运算符，例如卷积或矩阵乘法能以非常不同的方式映射和优化在不同的硬件后端。这些特定硬件优化在内存布局、并行线程模式、缓存访问模式和硬件基元的选择方面，可能会发生巨大的变化。我们希望能以通用方式对此进行明确表达。

我们建立了一个低级表示来解决这个问题。这个表示基于索引公式，而且支持重复计算。

低级别IR采用了现有的图像处理语言（例如Halide或darkroom）的原理来制定一种表现力很强的深度学习DSL。TVM在循环变换工具（例如loopy等）的启发下构建了图优化。我们也从MXNet、TensorFlow、Theano等深度学习框架的数据流描述语言中获得灵感。然后在调度阶段对TVM中描述的算法进行处理，以应用针对目标硬件后端订制的转换。

TVM包括CPU优化框架中常见的标准转换基元。更重要的是，TVM集成了针对GPU的新优化基元，包括利用线程协作模式、数据布局变换和强大的计算基元。将TVM和NNVM结合使用，可以用多种方式优化软件堆栈中的深度学习工作负载，进一步实现计算图级和运算符级的优化。

多语言和平台支持

TVM的优势之一，就是对多个平台和语言提供了丰富的支持。这由两个部分组成。一是编译器堆栈，其中包括完整的优化库，以产生优化过的机器代码；二是轻量级的运行环境，提供了在不同平台上部署编译模块所需的可移植性。

TVM目前支持嵌入式编译器堆栈的Python和C++接口。我们在设计框架时最大程度的实现了重复利用，以便编译器堆栈的改进可以在Python和C++组建之间互换使用。

我们还提供了一个轻量级的运行环境，可以让TVM用JavaScript、Java、Python、C++等编译过的代码，运行在Android、iOS、树莓派和网页浏览器等平台上。

远程部署和执行

通过轻量级的接口TVM RPC，可以在远程嵌入式设备上部署和执行TVM交叉编译过的模块。这为TVM用户提供了一个熟悉的高级Python界面，用于在各种低级嵌入式设备上远程编译、优化和测试深度学习算法。

性能

TVM处于早期阶段，让然有很多需要改进的地方，但是我们已经看到一些令人振奋的结果。

树莓派

我们首先在一个树莓派3B上运行ResNet工作复杂，来比较TVM和nnpack。由于时间限制，我们使用了TVM实现直接卷积，而nnpack用于对3×3内核执行winograd卷积。

我们发现在这次的树莓派试验中，使用TVM的自动调整内核，我们可以获得类似nnpack中的手动优化内核的性能。

GPU

以下测试归功于Leyuan Wang (AWS / UCDavis)，Yuwei Hu(TuSimple)，Weitang Liu (AWS/ UCDavis)。

作为一个概念的证明，我们创建了一个端到端的编译流程，可以将MXNet模型编译成TVM执行图。我们通过自动融合运算符并使TVM生成融合的内核，在图形节点之间应用优化。我们对MobileNet ImageNet工作负载进行了测试，结果如下。

结果显示，TVM在速度方面胜过我们的基准算法。更有趣的是，内核融合带来额外的加速。值得一提的是，TVM自己能生成所有优化的GPU内核，耳部依赖于CuDNN等外部库。

我们正在开展更多的实验，并准备放出更新的结果。

源代码

GitHub地址在此：

https://github.com/dmlc/tvm

— 完 —

加入社群

量子位AI社群7群开始招募啦，欢迎对AI感兴趣的同学，加小助手微信qbitbot2入群；

此外，量子位专业细分群(自动驾驶、CV、NLP、机器学习等)正在招募，面向正在从事相关领域的工程师及研究人员。

进群请加小助手微信号qbitbot2，并务必备注相应群的关键词~通过审核后我们将邀请进群。（专业群审核较严，敬请谅解）

诚挚招聘

量子位正在招募编辑/记者，工作地点在北京中关村。期待有才气、有热情的同学加入我们！相关细节，请在量子位公众号(QbitAI)对话界面，回复“招聘”两个字。

量子位 QbitAI

վ'ᴗ' ի 追踪AI技术和产品新动态