专栏 | 陈天奇等人提出AutoTVM:让AI来编译优化AI系统底层算子

选自:TVM社区

作者:陈天奇

Alexnet 横空出世的两年之前,我选择了深度学习作为本科毕设方向。当时的深度学习并没有现在那么火热,也没有现有的 GPU 卷积实现,我在实验室的 GTX 470 上面第一次手写 GPU 卷积用以支持卷积 RBM,当时针对 CPU 十多倍的加速比让我兴奋不已。虽然最终那段研究经历并没有结果,但是计算本身对于机器学习的推动作用却深深地印在了我的脑海里。

深度学习系统支是推动进步的引擎,而在引擎的核心,则是像 cuDNN 这样的针对硬件的高效算子实现。每一个高效的算子库的背后都是大量工程的工程努力。虽然我自己已经从事了几年深度学习系统的工作,如何实现高效的硬件代码始终对于我来说是一个头疼的问题。最近的深度学习系统趋势开始进行更多的高层优化:包括探索从不同的数据表示,精度和算子融合等各个方面,而这些优化势必会带来更多的算子版本需要去优化,这无疑是一个非常大的工程问题。我们还需要把这个问题乘上包括移动芯片,移动 GPU 和 AI 加速器等等后端硬件的种类。针对每一个平台都去实现一遍底层的优化代价无疑会非常大。

我们希望利用机器学习本身去解决这个问题,以深度学习编译器 TVM 为基础自动化地优化深度学习算子性能。写到这里,第一个大家会问的问题是为什么机器有可能会比手工做的优化要好。其实具体原因也很简单,一个人的精力是有限的,我们可以针对几个特定的场景去做一些优化。而机器有无限的精力可以针对每一种类型,数据表示,算子融合和实际的输入尺寸做特定的调整。像算子融合,这样的优化必须要利用自动代码生成才可以做到。所以机器在一定程度上可以取巧。

如何可以让机器匹配手写优化性能

思路非常直接:

  1. 建立一个足够大的搜索空间,保证可能的人工手写优化全部包含在这个搜索空间里面
  2. 快速地搜索这个这个空间,获取优化的实现

这两个想法很简单,其实在两年之前我们基本上就有了这一个解决问题的蓝图,但是执行都非常困难。第一个问题是如何建立足够大的搜索空间,使得其可以包含人可能达到的手写优化。这个在之前 TVM 的介绍里面我们提到过,我们通过参考改进 Halide 的 schedule 编程源语,加入了对于 GPU,加速器的抽象支持,通过一年的努力使得搜索空间基本可以匹配手写优化的性能。值得指出的是手工的优化可以非常多样性,并不是所有的优化都可以被简单的一个固定 pattern 总结。这也是之前很多自动代码生成框架无法达到比较好的效果的一个原因。我们也需要不断地总结抽象的优化规律,加入到 TVM 的原语中。另外在必要的时候我们也可以通过 tensorize 这个抽象可以插入一些手工优化的 micro kernel,但是依然使用框架来做外层调度,达到最好的效果。

有个足够好的搜索空间,剩下的问题是如何在几十亿的可能性里面去选择比较好的实现。这里有几种常见的做法。传统的高性能计算库会采用 Auto Tuning,也就是把可能的参数都尝试一遍。这样做的潜在问题是空间太大之后枚举开销过高。另外一种常见的做法是类似于数据库的做法,针对程序建立一个代价估价函数,然后利用估价函数来搜索。这个做法可能碰到的主要问题是估价函数不一定可以估计准确,并且针对每个新的硬件特性必须要重新设计估价函数。

我们的利用机器学习来学习程序空间的代价估价函数。具体地说,探索程序在一开始会随机地选取一些设定,直接到硬件上面去运行生成的代码,再通过得到的反馈数据来更新我们的程序代价估计函数。这里面比较有趣的一点是模型的可迁移性。因为真正的深度学习系统需要优化许多不一样输入类型,输入形状的算子。一个可迁移的模型可以通过学习已经看到过的算子优化记录来预测新的目标的代价,导致最后的搜索时间可以大幅降低。

其实在这个整个优化的过程中机器学习算法只是起到了重要但是只是一部分的作用。在一些如 ARM 我们的先验模版比较强的情况下,随机搜索就可以达到非常不错的效果。但是重要的是我们需要建立集群化,自动化,python 可以对接的框架来做这些实验。这里特别提一下 TVM 最近引入的 RPC tracker,支持了这一个功能,才可以使得我们可以在多种设备上进行实验。

论文链接

对于具体实验效果和方法介绍有兴趣的同学可以看论文里面的结果

  • 机器学习:https://arxiv.org/abs/1805.08166
  • 系统: https://arxiv.org/abs/1802.04799

未来展望

在可见的未来里面我们还是需要高性能优化的知识不断充实完善 TVM。手写算子优化和自动优化可能会在比较长的一段时并存下去。但是我们个人相信更多的自动化和半自动话也会成为主流。文章中提到的搜索空间,RPC 集群支持都已经正式加入到 TVM 中去。自动化模块也会在下一个版本和大家见面。

围绕这个问题,这个方向还留下包括模型表示,迁移等在内还有很多未知的问题等待着大家去探索,也欢迎更多的小伙伴加入 TVM 社区。

后记

后续我们也会继续在这一专栏介绍 TVM 深度学习自动优化等相关方向的内容,欢迎有兴趣的同学投稿。包括本次工作在内的 TVM 工作是整个社区团队一起努力的结果(而并非第一作者作者个人),转载本专栏文章时请注明 TVM 社区。

本次机器学习相关工作主要成员来自华盛顿大学,上海交大和复旦大学。我本人十分感谢交大 acm 学习经历对于这次工作的影响:十年之前的这个时候我在交大 acm 班正好在做编译原理大作业。本次论文的主力之一,第二作者郑怜悯也来自交大,是今年编译原理的助教。我们当时编译原理课程主现在依然挂在李沐在交大的主页上 http://bcmi.sjtu.edu.cn/home/limu/tiger/ 。

原文发布于微信公众号 - 机器之心(almosthuman2014)

原文发表时间:2018-05-23

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

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏灯塔大数据

深度|DT时代的核心竞争力---数据分析与挖掘

数据分析与挖掘,指的是通过对大量的数据进行观察与分析。发掘其中的未知的,潜在的、对决策有价值的关系、模式和趋势,并利用这些规则建立决策模型、提供预测性支持的方...

3044
来自专栏机器之心

人工智能与自然语言处理概述:AI三大阶段、NLP关键应用领域

选自xenonstack 机器之心编译 参与:黄小天、李亚洲 近日,Xenonstack 上推出了一篇名为《Overview of Artificial In...

3978
来自专栏大数据文摘

如何把各类难题变得数据可解?Get与数据科学家聊天的正确姿势

834
来自专栏AI科技大本营的专栏

AI 技术讲座精选: 数学不好,也可以学习人工智能(二)

【AI100 导读】本系列文章将陆续向大家推荐一些数学用书,今天这篇文章有针对性的介绍了数学不好的人,究竟该怎样学习人工智能。 ? 如果你已经看过本系列的第一篇...

3516
来自专栏每周一脱topic

推荐系统-学习总结

推荐系统目前几乎无处不在,主流的app都基本应用到了推荐系统。

1.1K13
来自专栏数据科学与人工智能

【深度学习】Github上的十大深度学习项目

本文作者Matthew May是一位正在进行并行式机器学习算法研究的计算机硕士研究生,同时Matthew也是一位数据挖掘研习者,数据发烧友,热忱的机器学习科学家...

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

干货分享:数据可视分析中的知识产生模型

可视分析技术已经发展了近十年。在这些年间,人们研究了大量的可视分析方法和案例,发表了不少研究论文。然而,对于一些基本问题,人们依然没有明确的答案。例如,一个基本...

3886
来自专栏AI科技评论

不一样的论文解读:2018 KDD best paper「Embeddings at Airbnb」

Airbnb 的 Real-time Personalization using Embeddings for Search Ranking at Airbnb...

652
来自专栏CDA数据分析师

专栏 | 案例:电信用户分群精准画像的7个步骤

“每天一个数据”分析师新一期内容奉上,请享用~ 转载请注明来自CDA数据分析师 否则小编将举报到底! 本期我们有幸采访到的嘉宾名叫兰锦池,2012年硕士毕业,...

2289
来自专栏FD的专栏

人工智能与自然语言处理概述:AI三大阶段、NLP关键应用领域

AI 指代「人工智能」,是让机器能够像人类一样完成智能任务的技术。AI 使用智能完成自动化任务。

783

扫码关注云+社区