Jeff Dean推荐:用TPU跑Julia程序,只需不到1000行代码

来源:arXiv 编辑:肖琴
【新智元导读】将Julia代码直接部署到谷歌Cloud TPU,让程序运行更快的官方指南来了!Julia和TPU的结合意味着快速、易于表达的ML计算!”

Julia是一门集众家所长的编程语言。随着Julia 1.0在8月初正式发布,Julia语言已然成为机器学习编程的新宠。

这门由 MIT CSAIL 实验室开发的编程语言结合了 C 语言的速度、Ruby 的灵活、Python 的通用性,以及其他各种语言的优势于一身,并且具有开源、简单易掌握的特点。

随着用户越来越多,围绕Julia的开发工具、技术、教程等也愈加丰富。昨天,Julia开发人员Keno Fischer和Elliot Saba发表了一篇新论文Automatic Full Compilation of Julia Programs and ML Modelsto Cloud TPUs,介绍如何将Julia代码直接部署到Google Cloud TPU,让程序运行更快

Jeff Dean在推特上推荐了这篇论文,评价称:“Julia和TPU的结合意味着快速、易于表达的ML计算!

谷歌的Cloud TPU是一种用于机器学习工作负载的很强大的新硬件架构。近年来,Cloud TPU为谷歌的许多里程碑式的机器学习成就提供了动力。

谷歌现在已经在他们的云平台上开放提供一般用途的TPU,并且最近已经进一步开放,允许非TensorFlow前端使用。

这篇论文描述了通过这个新的API和Google XLA编译器,将Julia程序的适当部分卸载(offload)到TPU的方法和实现。

这一方法能够将表示为Julia程序的VGG19模型的前向传递(forward pass)完全融合到单个TPU可执行文件中,以便卸载到设备。该方法也很好地与Julia代码上现有的基于编译器的自动微分技术相结合,因此我们也能够自动获得VGG19的反向传递并类似地将其卸载到TPU。

使用这一编译器定位TPU,能够在0.23秒内对100张图像的VGG19前向传递进行评估,这与CPU上原始模型所需的52.4秒相比大幅加速了。他们的实现仅需不到1000行Julia代码,没有对核心Julia编译器或任何其他Julia包进行TPU特定的更改。

具体方法和实现细节请阅读原论文。以下主要从分别从回顾TPU硬件架构、Julia编译器的workflow、将XLA嵌入到Julia IR的细节,以及结果与讨论几个部分进行介绍。

谷歌TPU和XLA编译器

2017年,谷歌宣布他们将通过云服务向公众提供其专有的张量处理单元(TPU)机器学习加速器。最初,TPU的使用仅限于使用谷歌的TensorFlow机器学习框架编写的应用程序。幸运的是,2018年9月,Google通过较低级别的XLA(Accelerated Linear Algebra)编译器的IR开放了对TPU的访问权限。该IR是通用的,是用于表示线性代数原语的任意计算的优化编译器,因此为非Tensorflow用户以及非机器学习工作负载的TPU目标提供了良好的基础。

XLA(加速线性代数)是谷歌的一个部分开源编译器项目。它具有丰富的输入IR,用于指定多线性代数计算,并为CPU,GPU和TPU提供后端代码生成功能。XLA的输入IR(称为HLO高级优化IR)在基本数据类型或其元组(但没有元组数组)的任意维数组上运行。HLO操作包括基本算术运算、特殊函数、广义线性代数运算、高级数组运算以及用于分布式计算的原语。XLA可以执行输入程序的语义简化,以及执行整个程序的内存调度,以便有效地使用和重用可用内存(这是大型机器学习模型的一个非常重要的考虑因素)。

每个HLO操作都有两种操作数:

  1. 静态操作数,它的值必须在编译时可用并配置操作。
  2. 动态操作数,由上述张量组成。

这篇论文介绍了使用这个接口将常规的Julia代码编译带TPU的初步工作。这一方法不依赖跟踪,而是利用Julia的静态分析和编译功能来编译完整的程序,包括对设备的任何控制flow。特别是,我们的方法允许用户在编写模型时充分利用Julia语言的完整表现力,能够编译使用Flux机器学习框架编写的完整机器学习模型,将前向和后向模型传递以及训练loop融合到单个可执行文件,并将其卸载到TPU。

Julia编译器的工作原理

为了理解如何将Julia代码编译为XLA代码,了解常规Julia编译器的工作原理是有益的。Julia在语义上是一种非常动态的语言。但是,在标准配置中,Julia的最终后端编译器是LLVM(Lattner&Adve,2004),它是一个静态编译器后端。

Julia编译器需要将语言的动态语义与LLVM表示的静态语义之间联系起来。为了理解这个过程,我们将研究Julia系统的四个方面:动态语义、静态编译器内部函数的嵌入、过程间类型推断,以及静态子图的提取。此外,我们还将研究这些特征与宏和生成的函数的交互,这些函数将与XLA编译器相关。

如何将XLA嵌入到Julia IR

XLA嵌入

要编译为XLA而不是LLVM,我们应用了上一节中概述的策略。实际上,我们可以重用大多数编译器本身(特别是所有类型推断和所有mid-level优化传递)。

让我们先定义动态语义和静态嵌入。

张量表示(Tensor representation)

由于其作为线性代数的教学和研究语言的传统,Julia具有非常丰富的数组抽象层次结构。Julia的标准库数组是可变的,并且在类型和维度上进行参数化。此外,StaticArrays.jl(Ferris&Contributors,2018)包提供了在元素类型和形状上进行参数化的不可变数组。因此,成形的N维不可变张量的概念对Julia代码来说并不陌生,并且大多数现有的通用代码能够毫无问题地处理它。

因此,我们通过定义一个runtime结构来嵌入XLA values。

Listing 1: XRTArray3的定义。

操作表示(Operation representation)

分离静态和动态操作数

HLO操作数(HLO operands)分为静态和动态操作数。假设我们有一个示例XLA操作'Foo'采用一个静态操作数(例如一个整数)和两个动态操作数。这个嵌入如下所示:

在这个示例中,“execute”函数实现在远程设备上运行操作的动态语义。函数(hlo::HloFoo)(...) 语法表示调用运算符重载。因此,这意味着对HloFoo(1) 的调用将构造并返回一个callabale对象,当在两个XRTArrays上调用时,它将使用静态操作数'1'远程执行'Foo'HLO操作,并且对应于两个数组的动态操作数。这种分离并不是绝对必要的,但确实有嵌入到Julia IR的有用特性,易于理解:

在Listing 2的示例中,我们将HLO操作数(包括静态操作数)拼接到AST中。这产生了一个非常简单的XLA映射(遍历每个语句,从拼接指令规范获取静态操作数,从类型推断获得动态形状并生成相应的XLA代码)。

当然,我们通常不会手动拼接这些指令,但是手动拼接的示例说明了为什么分离静态操作数很有用,并说明了成功offload到XLA的条件。

如果经过所有相关的Julia级别优化之后,IR可以完全卸载:

Listing 2: 手动构建的XLA嵌入

满足这些条件的IR可以简单地转换成XLA IR。

结果

本文描述的方法在很大程度上依赖于Julia中间端编译器,以确定足够精确的信息,在程序的足够大的子区域中分摊任何启动开销。

在本节中,我们证明了Julia编译器确实足够精确,使该方法适用于实际的程序。

VGG19 forward pass

图1:在编译到XLA之后,Metalhead.jl VGG19的forward pass 和backwards pass 生成的XLA指令摘要。

这里显示了未优化(在Julia前端之后)和优化的计数(在类似于CPU后端使用的XLA优化pipeline之后,但没有HLO融合)。

VGG19 backward pass

为了获得backwards pass,我们使用基于Zygote.jl编译器的AD框架(Innes, 2018)。Zygote对Julia代码进行操作,其输出也是Julia函数(适合重新引入Zygote以获得更高阶导数,也适合编译到TPU)。

示例如下:

结论

在这篇论文中,我们讨论了如何将Julia代码编译为XLA IR,从而实现卸载到TPU设备。这里描述的实现重新利用了现有Julia编译器的重要部分,因此所有代码不到1000行,但是仍然能够编译模型的forward和backward pass(及其融合,包括 training loop)到单个XLA内核,模型例如VGG19。

我们还演示了Julia的多重调度语义如何在这个转换的规范中提供帮助。这项工作表明,不仅可以将用Julia编写的多个ML模型编译到TPU,而且可以编写更通用的非ML Julia代码(只要这些代码也由线性代数操作控制)。我们希望这可以加速对非ML问题领域的探索,TPU可能对这些领域有用。

论文地址:

https://arxiv.org/pdf/1810.09868.pdf

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

原文发表时间:2018-10-25

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

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏数据魔术师

运筹学教学|修正单纯形法(revised simplex algorithm)代码分享及详细注释

欢声笑语中,小编学会了单纯形法,心里还有点小傲骄!!准备晚上去PUBG里面潇洒一把~ ? 然而,老板突然来电话说,单纯形法有升级的版本!需要我赶紧准备一份代码。...

8817
来自专栏架构师之路

洗稿,技术上怎么判断文章相似性?

这几天“差评洗稿”的事情,闹得沸沸扬扬,楼主本身也是内容创作的手艺人,简单说两句。

1964
来自专栏逍遥剑客的游戏开发

PhysX学习笔记(3): 动力学(2) Actor

1982
来自专栏宏伦工作室

基于维基百科的中文词语相关度计算

1953
来自专栏素质云笔记

R︱Yandex的梯度提升CatBoost 算法(官方述:超越XGBoost/lightGBM/h2o)

俄罗斯搜索巨头 Yandex 昨日宣布开源 CatBoost ,这是一种支持类别特征,基于梯度提升决策树的机器学习方法。 CatBoost 是由 Yand...

5709
来自专栏机器之心

资源 | 知网(HowNet)知识库的简单调用指南

机器之心整理 参与:蒋思源 机器之心曾采访过语知科技的董强先生,在那一篇文章中,我们详细讨论了基于知网知识库的 NLP 解决方案。虽然我们已经了解了这种方法的潜...

7155
来自专栏落影的专栏

OpenGL ES实践教程(九)OpenGL与视频混合

前言 前面的实践教程: OpenGL ES实践教程1-Demo01-AVPlayer OpenGL ES实践教程2-Demo02-摄像头采集数据和渲染 O...

5695
来自专栏Aloys的开发之路

OOAD与UML笔记

UML基础介绍 1.UML的定义 统一建模语言(UML)是一种图形化的语言,它可以帮助我们在OOAD过程中标识元素、构建模块、分析过程并可通过文档说明系统中的重...

1998
来自专栏好好学java的技术栈

“365算法每日学计划”:03打卡-贪心算法

自从开始做公众号开始,就一直在思考,怎么把算法的训练做好,因为思海同学在算法这方面的掌握确实还不够。因此,我现在想做一个“365算法每日学计划”。

2382
来自专栏斑斓

MongoDB的数据建模

MongoDB是一种面向Document的NoSQL数据库,如果我们还是按照RDB的方式来思考MongoDB的数据建模,则不能有效地利用MongoDB的优势;然...

3446

扫码关注云+社区

领取腾讯云代金券