专栏 | 陈天奇等人提出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 条评论
登录 后参与评论

相关文章

来自专栏机器之心

教程 | 教Alexa看懂手语,不说话也能控制语音助手

数月前的某个夜晚,我躺在床上时,一个念头闪过我的脑海——「如果语音是计算接口的未来,那么那些听不见或看不见的人该怎么办?」我不知道究竟是什么触发了这个想法。我自...

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

Python实现基于SVM的股票预测

核心 因为是分类算法,因此不像ARIMA一样预测的是时序。分类就要有东西可分,因此将当日涨记为1,跌记为0,作为分类的依据。使用历史数据作为训练数据。 处理数据...

4585
来自专栏AI研习社

最知名的5个机器学习框架,知道一个说明你入门了

AI 研习社消息:业内知名数据科学网站 KDnuggests,昨日评选出了四月份“你不可忽视的五个机器学习项目”。 你可能没听过它们,但今天或许会考虑上手。至于...

4289
来自专栏生信宝典

数据可视化基本套路总结

真依然很拉风,简书《数据可视化》专栏维护者,里面有很多优秀的文章,本文便是其中一篇。

1052
来自专栏AI研习社

想在手机上用自己的 ML 模型?谷歌模型压缩包你满意

AI 研习社按: Google I/O 2018 上,谷歌发布了可供开发者定制移动端机器学习模型的 ML Kit 开发套件,关于该套件中的核心技术:Learn2...

691
来自专栏专知

ACL 2018教程:关于语义和语用的100件事(附下载)

【导读】近日,ACL 2018 在澳大利亚墨尔本举办,举办地点为墨尔本会展中心。其tutorial 旨在帮助领域新手了解计算机语言学与自然语言处理的最新进展以及...

850
来自专栏牛肉圆粉不加葱

CTR 预估候选集选择方案

最近一段日子,怎么在一个大的新闻候选集合(假设有10w条新闻)选出用来做 CTR 预估的小集合(假设是100条新闻)困扰我们已久,昨夜我夜观星象,突来灵感,遂有...

862
来自专栏大数据文摘

手把手|用Python端对端数据分析识别机器人“僵尸粉”

1646
来自专栏量子位

大手笔!YouTube放出史上最大已标注视频数据集

以下内容来自Google Research Blog,量子位编译 ? 当今机器学习中最具挑战性的研究领域之一,是使计算机能够理解一个场景是什么。例如,虽然人类知...

18110
来自专栏量子位

用深度学习拯救手抖星人!Facebook详解全景照片修复技巧

陈桦 编译整理 量子位 报道 | 公众号 QbitAI 拍全景照片,重要的是手要稳,手要稳,手要稳……或者支个三角架。 ? 上面视频中这位小姐姐的水平,可以说是...

3077

扫码关注云+社区