陈天奇团队发布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

原文发布于微信公众号 - 量子位(QbitAI)

原文发表时间:2017-08-18

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏腾讯玄武实验室的专栏

让 iPhone “崩溃” 又有了新方法:只需要一个视频

玄武实验室在 10 月 15 日就发现了视频样本,在深入分析后,我们在微信后台对能够触发这种漏洞的恶意视频进行检测和拦截,保护了大量用户免遭攻击。

5722
来自专栏编程

Python数据分析系列(2)——美国纽约皇后区空气质量分析

作者:王大伟 Python爱好者社区唯一小编 博客:https://ask.hellobi.com/blog/wangdawei 关注Python爱好者社区回复...

2955
来自专栏机器之心

业界 | 英伟达开源硬件加速项目NVDLA:一种标准化的推断加速框架

3509
来自专栏AI研习社

博客 | Github开源人体姿态识别项目OpenPose中文文档

人体姿态识别与估计的应用场景:抖音尬舞机、体育动作教学、3D健身教练、3D试衣、绘画辅助、游戏人物动作采集。

1.1K4
来自专栏AI派

还没准备好数据呢,为什么要着急用算法呢

开始之前,通知下我的读者,随着订阅读者越来越多,为了对读者们负责,有以下几件事情需要告知下:

2176
来自专栏ATYUN订阅号

2017最流行的十大Python库

年末是你需要思考过去一年的成就的时候。对于程序员来说,这通常是回顾今年发布的或者最近流行的开源库,因为它们是解决特定问题的绝佳工具。 在过去的两个多年中,我们都...

4247
来自专栏快乐八哥

数据可视化-EChart2.0使用总结1

图表是企业级Web开发必不可少的一个功能点。也是“数据可视化的一个具体呈现”。今天看到阮一峰翻译的“数据可视化:基本图表”一文,同时梳理一下公司现在项目使用的E...

2905
来自专栏流柯技术学院

性能测试学习之三—— PV->TPS转换模型&TPS波动模型

TPS = ( (80%*总PV)/(24*60*60*(T/24)))/服务器数量

2332
来自专栏ios 技术积累

ios 百度地图设置BMKAnnotationView层级

产品有一个新需求,就是百度地图上显示的BMKAnnotationView,根据数值越大的放在最上层,原因是因为BMKAnnotationView很多并且有重叠的...

3134
来自专栏机器之心

自动生成硬件优化内核:陈天奇等人发布深度学习编译器TVM

3606

扫码关注云+社区

领取腾讯云代金券