【深度学习框架的未来】李沐:AWS开源端到端AI框架编译器NNVM

【新智元导读】AWS AI主任科学家李沐今天在官博发表文章,介绍了一种新的编译器NNVM,可以大大简化新的AI前端框架和后端硬件的设计,在各种前端和后端为用户提供一致的结果。NNVM编译器将高级计算图编译成优化的机器代码,只需花费很少的精力,就可以匹配甚至超过两个完全不同的硬件(ARM CPU和Nvidia GPU)上的state-of-the-art性能。

有很多个人工智能(AI)框架供你选择来开发AI算法。同样你也可以选择各种硬件来训练和部署AI模型。框架和硬件的多样性对于保持AI生态系统的健康至关重要。但是,这种多样性也为AI开发人员带来了几个挑战。本文将简要介绍这些挑战,并引入一个编译器解决方案来辅助解决这些问题。

下面,我们首先回顾挑战,介绍华盛顿大学和AWS研究团队,然后介绍编译器的工作原理。

三大挑战

首先,由于前端接口和后端实现之间的差异,从一个AI框架切换到另一个AI框架是很麻烦的。此外,算法开发人员在开发和交付流程中可能会使用多个框架。在AWS,我们有客户希望在MXNet上部署其Caffe模型,从而享受Amazon EC2上的加速性能。根据Joaquin Candela最近的博客,用户可能会使用PyTorch快速开发,然后在Caffe2上部署。但是,我们也听到有人表示在将模型从一个框架转换到另一个框架之后调试结果差异十分困难。

其次,框架开发人员需要维护多个后端,才能确保从智能手机芯片到数据中心GPU的各种硬件的性能。以MXNet为例,MXNet有一个从头构建的便携式C++实现,还附带了专门的后端支持,比如英伟达GPU的CuNNN和英特尔CPU的MKLML。确保这些不同的后端为用户提供一致的数值结果十分具有挑战性。

最后,芯片厂商制造的每款新芯片也需要支持多个AI框架。每个框架中的workload都以独特的方式进行表征和执行,因此即使像卷积一样的运算也可能需要以不同的方式进行定义。因此,在芯片设计和制造方面,支持多个框架也需要大量的工程设计。

华盛顿大学和AWS研究团队

不同的AI框架和硬件为用户带来巨大的收益,但AI开发人员为终端用户提供一致性却是非常具有挑战性的事情。幸运的是,我们并不是第一个面对这种问题的人。计算机科学在不同硬件上运行各种编程语言的历史由来已久。解决这个问题的一个关键技术是编译器。由编译技术驱动,来自华盛顿大学(UW)计算机科学与工程系保罗·艾伦学院的陈天奇,Thierry Moreau,Haichen Shen,Luis Ceze,Carlos Guestrin和Arvind Krishnamurth等一批研究人员,与AWS AI团队的Ziheng Jiang,提出了TVM堆栈来简化这个问题。

今天,AWS很高兴与UW研究团队一起宣布基于TVM堆栈的端到端编译器,它将workload直接从各种深度学习前端编译成优化的机器代码。

架构

下面,我们先来看这个编译器架构。

我们可以注意到,典型的AI框架大致分为三部分:

  • 前端开放一个易于使用的界面给用户;
  • 从前端收到的workload通常被表示为由数据变量(a,b和c)和运算符(*和+)组成的计算图;
  • 从基本的算术运算到神经网络层的运算符都可以针对多个硬件进行实现和优化。

新的编译器,NNVM编译器,基于TVM堆栈中的两个组件:用于计算图的NNVM和用于张量运算符的TVM。

NNVM - 计算图中间表示(IR)堆栈

NNVM的目标是将不同框架的workload表示为标准计算图,然后将这些高级图转换为执行图。这种计算图的灵感来自Keras中的层定义和numpy的张量运算符。

NNVM还随附例程,名叫Pass,遵循LLVM约定。这些例程可以向图中添加新属性来执行它们或修改图以提高效率。

TVM - 张量IR堆叠

源自Halide的TVM实现了计算图中使用的运算符,并针对目标后端硬件进行了优化。与NNVM不同,TVM提供了一种独立于硬件的域特定语言,简化张量索引级别中的操作符实现。TVM还提供调度原语(如多线程,平铺和高速缓存)来优化计算以充分利用硬件资源。这些计划与硬件有关,可以手动编码,也可以自动搜索优化的模式。

支持的前端框架和后端硬件

支持的前端和后端如下图所示。

MXNet直接将其计算图转换为NNVM图得到支持,Keras也以类似的方式得到支持,不过正在开发中。NNVM编译器也可以采用模型格式,如CoreML。因此,凡是能够使用这些格式框架都可以使用此编译堆栈。

TVM目前随附多个代码生成器,支持各种后端硬件。例如,TVM为CPU(如X86和ARM)生成LLVM IR,也可以为各种GPU输出CUDA,OpenCL和Metal内核。

添加新的支持的操作也很简单。对于新的前端,我们只需要将其workload转换为定义计算图和运算符规范的NNVM即可。要添加新的硬件,我们可以重复使用TVM的运算符实现,只需要指定有效schedule即可。

性能结果

我们使用MXNet展示NNVM编译器性能,前端硬件配置是:ARM上的Raspberry PI和AWS上的英伟达GPU。虽然这两种芯片间的架构差异巨大,但代码上的唯一差异在于调度部分。

GPU的schedule主要由Leyuan Wang(AWS)和Yuwei Hu(TuSimple)在实习期间编写。比较NNVM编译器与MXNet,英伟达K80的cuDNN后端。像深度卷积的运算符这样在cuDNN中没有有效支持的,通过手动优化的CUDA内核来实现。可以看出,NNVM编译器运行ResNet18和MobileNet比cuDNN后端略好(1.2倍快)。

在Raspberry PI的情况下,我们通过自动调谐器选择最佳schedule。我们通过在Raspberry Pi上对运算符性能进行基准测试,给出了每个运算符在给定形状的最佳schedule。

还是比较NNVM编译器与MXNet。MXNet默认启用OpenBLAS和NNPACK,我们还手动打开了NNPACK中的winograd卷积以获得最佳性能。

可以看出,NNVM编译器在Resnet18上快了2.2倍。而MobileNet上更有11.5倍的差异。这主要是由于在MXNet中没有优化深度卷积(因为dnn库中缺少这样的运算符),而NNVM编译器则受益于直接生成高效的代码。

结论

我们介绍了NNVM编译器,它将高级计算图编译成优化的机器代码。NNVM编译器基于TVM堆栈中的两个组件:NNVM使用图优化例程提供计算图和运算符的specification,运算符通过使用TVM针对目标硬件实现和优化。我们证明,只需花费很少的精力,这种编译器可以匹配甚至超过两个完全不同的硬件(ARM CPU和Nvidia GPU)上的state-of-the-art性能。

我们希望NNVM编译器可以大大简化新的AI前端框架和后端硬件的设计,并帮助在各种前端和后端为用户提供一致的结果。

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

原文发布于微信公众号 - 新智元(AI_era)

原文发表时间:2017-10-07

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

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏生信小驿站

Python Vs R:数据科学家的永恒问题pythonR结论

Python有一些使用案例,R也是如此。使用它们的场景各不相同。 更常见的是环境以及客户或雇主的需求决定了Python和R之间的选择。许多事情在Python中都...

912
来自专栏大数据挖掘DT机器学习

【推荐】分析的前提—数据质量

数据质量(Data Quality)是数据分析结论有效性和准确性的基础也是最重要的前提和保障。数据质量保证(Data Quality Assurance...

3115
来自专栏wblearn

UML学习归纳整理

之前在学校比较系统的学习过统一建模语言UML,但长时间没使用遗忘了许多,最近因工作需要,所以对UML重新学习和梳理一下。UML可帮助我们做软件需求分析和软件设计...

971
来自专栏哲学驱动设计

重构实践:体验interface的威力(一)

背景     GIX4是一个建筑行业的指标计算软件,用于数据统计、分析。导入的大量数据,大部分呈现逻辑上的树状结构(关于它的重构,见:《重构一个繁琐的数据结构》...

1977
来自专栏PPV课数据科学社区

【连载】R语言18讲(四)

? 前面讲到了怎么导入数据,获取数据,按照数据挖掘的流程,在我们导入数据之后接下来的工作便是对数据进行初步的探索,探索的工作主要包括两大块,第一是数据...

2868
来自专栏AI研习社

谷歌云工程师亲自示范:新发布的 Video Intelligence API 究竟怎么用

AI研习社按:3 月初,谷歌在 Google Cloud Next 2017 大会上对外公布了 Cloud Machine Learning API 更新。此次...

3337
来自专栏PPV课数据科学社区

做数据分析必须学R的4个理由

R 是一种灵活的编程语言,专为促进探索性数据分析、经典统计学测试和高级图形学而设计。R 拥有丰富的、仍在不断扩大的数据包库,处于统计学、数据分析和数据挖掘发展的...

2786
来自专栏生信技能树

一篇文章学会miRNA-seq分析

第一讲:文献选择与解读 前阵子逛BioStar论坛的时候看到了一个关于miRNA分析的问题,提问者从NCBI的SRA中下载文献提供的原始数据,然后处理的时候出现...

1.6K7
来自专栏令仔很忙

软件工程---软件需求分析

   需求分析是软件定义时期的最后一个阶段,它的基本任务是准确回答“系统必须做什么?”

1222
来自专栏华章科技

做数据分析必须学R的4个理由

R 是一种灵活的编程语言,专为促进探索性数据分析、经典统计学测试和高级图形学而设计。R 拥有丰富的、仍在不断扩大的数据包库,处于统计学、数据分析和数据挖掘发展的...

591

扫码关注云+社区