会员服务 ·

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

2017 年 8 月 18 日 量子位 专注报道AI

允中编译整理
量子位出品 | 公众号 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中描述的算法进行处理，以应用针对目标硬件后端订制的转换。