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

简介:
本文来自AI新媒体量子位(QbitAI)

f0e4bbe9643d5ac559115bf5a2129832f555f172

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

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

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

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

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

随后陈天奇也补充说:

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

非常建议大家尝试一下。

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

9fdc3e37339dbaa6188aab1e161b897f9cf64f55

关于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)堆栈,用来弥合深度学习框架和硬件后端之间的距离。

429011300e5015e4f12edaca9ce87af06de651dc

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

  • 为CPU、GPU和其他专用硬件,表示和优化常见的深度学习计算工作负载
  • 自动转换计算图以最小化内存占用,优化数据布局和融合计算模式
  • 提供端到端编译,从现有的前端框架到裸机硬件,直到浏览器可执行的javascript

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

a42d16bdc441722f705e321a0fd223b9c93bcb7f

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

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

技术细节

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

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

7b69367d875fcab7fcc3cdab0db3341819e9f6ec

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

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

62bb103240d37236224bc15c213f811f88df15f5

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

dc1b401f05e1ab69b74d20327240b8aa2418a336

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

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

eb33d613bab2f0d8e0e4748988707c123f9bed24

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

07ba6145db0fa6a12c96c44c45db12762cea6d45

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

多语言和平台支持

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

9e5220f734f930ba93ddd777fdf197dc622a0e94

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

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

远程部署和执行

e4f657b693ac0898109a92e78b7680f22df0f9af

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

性能

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

树莓派

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

30270562ea59a8dc2d3c521c5d36e590fe6023c5

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

GPU

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

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

4c18e9bc2eb73f8d5650ee66bb3ccc7669ea1679

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

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

源代码

GitHub地址在此:

https://github.com/dmlc/tvm

本文作者:允中
原文发布时间:2017-08-18 
相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
相关文章
|
3月前
|
机器学习/深度学习 边缘计算 PyTorch
PyTorch 与边缘计算:将深度学习模型部署到嵌入式设备
【8月更文第29天】随着物联网技术的发展,越来越多的数据处理任务开始在边缘设备上执行,以减少网络延迟、降低带宽成本并提高隐私保护水平。PyTorch 是一个广泛使用的深度学习框架,它不仅支持高效的模型训练,还提供了多种工具帮助开发者将模型部署到边缘设备。本文将探讨如何将PyTorch模型高效地部署到嵌入式设备上,并通过一个具体的示例来展示整个流程。
437 1
|
3月前
|
机器学习/深度学习 人工智能 前端开发
BladeDISC 深度学习编译器问题之BladeDISC在新硬件支持方面如何解决
BladeDISC 深度学习编译器问题之BladeDISC在新硬件支持方面如何解决
|
3月前
|
机器学习/深度学习 人工智能 前端开发
BladeDISC 深度学习编译器问题之实现硬件适配如何解决
BladeDISC 深度学习编译器问题之实现硬件适配如何解决
|
3月前
|
机器学习/深度学习 人工智能 自然语言处理
什么是深度学习模型?如何部署它?
【8月更文挑战第23天】
138 0
|
3月前
|
机器学习/深度学习 存储 物联网
深度学习模型的优化与部署
【8月更文第18天】随着深度学习技术的发展,模型规模变得越来越大,这对计算资源的要求也越来越高。为了能够在资源有限的边缘设备(如智能手机、物联网设备)上运行复杂的深度学习模型,我们需要采用一系列优化方法来减少模型大小和计算复杂度。本文将介绍几种常用的模型优化技术,并讨论如何在边缘设备或云端服务器上部署这些优化后的模型。
124 0
|
4月前
|
机器学习/深度学习 物联网 TensorFlow
使用Python实现深度学习模型:在嵌入式设备上的部署
【7月更文挑战第11天】 使用Python实现深度学习模型:在嵌入式设备上的部署
491 2
|
4月前
|
机器学习/深度学习 TensorFlow 算法框架/工具
使用Python实现深度学习模型:跨平台模型移植与部署
【7月更文挑战第10天】 使用Python实现深度学习模型:跨平台模型移植与部署
214 1
|
4月前
|
JavaScript Java 测试技术
基于springboot+vue.js+uniapp的多功能智能手机阅读APP附带文章源码部署视频讲解等
基于springboot+vue.js+uniapp的多功能智能手机阅读APP附带文章源码部署视频讲解等
75 1
|
3月前
|
机器学习/深度学习 存储 并行计算
深度学习之适应硬件的神经网络
深度学习的适应硬件的神经网络设计旨在最大限度地利用特定硬件平台的计算和存储能力,提高模型的执行效率和性能。这些硬件包括图形处理单元(GPU)、张量处理单元(TPU)、现场可编程门阵列(FPGA)和专用集成电路(ASIC)。
55 0
|
4月前
|
机器学习/深度学习 API Docker
使用Python实现深度学习模型:模型部署与生产环境应用
【7月更文挑战第7天】 使用Python实现深度学习模型:模型部署与生产环境应用
198 0