谷歌云TPU上可以用Julia啦 0.23秒跑100张图片 JeffDean点赞推荐

伊瓢 发自 凹非寺

量子位 报道 | 公众号 QbitAI

不久前，Julia Computing官方放出了一篇论文，展示将Julia代码和机器学习模型编译到谷歌云TPU的方法，可以实现在0.23秒内完成100张图片VGG19正向传递。

Jeff Dean看到这个方法后，开心的转了一作的Twitter，并且评价：真是又快又容易呀！

该方法应用的正是谷歌上个月给TPU开放许可权的XLA（Accelerated Linear Algebra）编译器。

XLA嵌入

首先，需要定义动态语义和静态嵌入，定义一个运行时结构来嵌入XLA值

XRTArray的定义

假设我们有一个示例XLA操作’Foo’采用一个静态操作数（例如一个整数）和两个动态操作数。 我们将声明此嵌入如下：

手动构建XLA嵌入：

将Julia映射到XLA

现在，可以将Julia代码编译到XLA，不过Julia不是用HLO运行的，而是根据Julia库提供的功能编写的。

好在，Julia使用多个调度可以很容易地表达如何用HLO实现标准库抽象。下面是一个简单的例子：

结果比对

论文中举了VGG19正向传递和反向传递的例子。

VGG19正向传递使用Metalhead软件包Mike Innes＆Contributors（2018）中的VGG19实现，它利用了Flux Innes＆Contributors（2017）框架将熟悉的机器学习层（卷积层，密集层）转换为线性代数运算。

但是，重要的是，Flux框架中的每一层都只是一个常规函数，而这些函数又调用常规线性代数运算。因此用Flux表达的机器学习模型只需要简单的Julia函数，适用于本文所述的方法。

VGG反向传递则使用基于Zygote.jl编译器的AD框架Innes（2018）。Zygote对Julia代码进行操作，其输出也是Julia函数，适合重新引入Zygote以获得更高阶导数，但也适合编译到TPU。

这张表格展示编译到XLA后，Metalhead.jl VGG19正向传递和反后传递的指令计数细分。

传送门

Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs

作者：Keno Fischer, Elliot Saba

https://arxiv.org/abs/1810.09868

—完—

加入社群

量子位AI社群开始招募啦，欢迎对AI感兴趣的同学，在量子位公众号（QbitAI）对话界面回复关键字“交流群”，获取入群方式；

此外，量子位专业细分群(自动驾驶、CV、NLP、机器学习等)正在招募，面向正在从事相关领域的工程师及研究人员。

进专业群请在量子位公众号（QbitAI）对话界面回复关键字“专业群”，获取入群方式。（专业群审核较严，敬请谅解）

活动策划招聘

量子位正在招聘活动策划，将负责不同领域维度的线上线下相关活动策划、执行。欢迎聪明靠谱的小伙伴加入，并希望你能有一些活动策划或运营的相关经验。相关细节，请在量子位公众号(QbitAI)对话界面，回复“招聘”两个字。