有了Julia语言,深度学习框架从此不需要计算图

选自julialang

作者:Mike Innes 等

机器之心编译

参与:刘晓坤、思源

本文基于 NeurIPS MLSys 的一篇论文《Fashionable Modelling with Flux》,探讨开发者们如何使用

Julia

语言从头开始思考机器学习工具,并提供对于现代机器学习工具所需改进的一些见解,涉及新的可微分编程工具 Flux、求梯度、支持 GPU 和 TPU、自动批处理。为什么 Julia 式的机器学习不需要计算图呢?因为 Julia 的机器学习语法就是计算图。

鉴于机器学习(ML)对编程语言、编译器和生态系统的众多需求,现在已经有很多有趣的发展。不仅 TensorFlow 和 PyTorch 等现有系统间的权衡得不到解决,而且这两个框架都包含不同的「静态图」和「eager execution」接口,但它们的形式已经比以前更加清晰。与此同时,机器学习模型基本上是可微分算法的思想(通常称为可微分编程)已经流行起来。

当前的机器学习框架遇到了阻碍,很多已有的新项目都完全移除了计算图,从而使可微分编程成为主流。例如,由 Theano 团队开发的 Myia 可以求微分并编译 Python 的一个子集为高性能 GPU 代码。Swift for TensorFlow 作为 Swift 语言的扩展,它可以将兼容的函数编译为 TensorFlow 计算图。最后,Flux 生态系统为 Julia 编译器提供了一些机器学习专用的工具,包括:first-class gradients、即时CUDA核编译、自动批处理(automatic batching)以及对新硬件(例如 TPU)的支持。

所有这些项目都有巨大的潜力,但目前看来 Julia 具有优势。

Flux 简介

我们需要一种语言来编写可微分算法,Flux 使 Julia 变成了这样的语言。Julia 专为数学和数值计算而设计,非常适合表达机器学习算法。同时,它在编译器中融合了现代设计和新思想,可以更轻松地满足尖端 ML 的高性能需求。

典型的框架通常包含数十万行 C++代码,Flux 却只有千行 Julia 代码。只需要一个求梯度的包(Zygote.jl)、一个用于 GPU 支持的包(CuArrays.jl)、再加上一些轻量函数,我们就能得到一个功能齐全的机器学习堆栈。

与其他下一代机器学习系统一样,Flux 致力于提供直观(「eager」或「define-by-run」)的接口,并对任何类型的计算图构建或性能注释进行严格控制。从控制流、数据结构到宏,Flux 支持语言的所有特征。用户可以在 Jupyter 笔记本中交互式地写代码,并将高性能数值计算与方便的绘图、可视化相结合。但我们也希望获得传统上由「静态图」框架所带来的好处,例如零开销源到源 AD、OP 融合、多 GPU /分布式训练和二进制部署等。

我们怎么能做到这一切?实际上,我们需要直接从 Julia 语法中提取和分析「静态图」,这实际完全上是编译器的正常工作。通过适当的角度来看,大多数机器学习系统问题都是标准的且经过充分研究的编译器问题。使用编译语言足以解决许多问题,扩展该编译器是解决更多问题的最佳方法。本文仅介绍了我们目前在该领域的工作范例,即求梯度、为 GPU 和 TPU 提供代码编译,以及自动批处理。

求梯度

推动反向模式求微分的极限,我们将此视为语言层面的问题。求微分是一种符号转换,属于编译器的领域。现有框架通过追踪(实际上是一种部分评估或抽象解释)来实现这一目标。人们引入了一种新的张量类型,它记录了所执行的所有基本数学运算,生成一个计算图(或符号表达式),其中删除了宿主语言的控制流和数据结构。然而,这给出了一个艰难的权衡:我们要么接受解释器的开销(eager execution),要么固定用户的控制流并限制可以构建的模型种类(静态图)。

反之,如果「计算图」就是 Julia 自己的语法呢?通过将这个想法发挥到极致,我们构建了 Zygote,它直接在 SSA 形式的中间表征(IR)上工作,支持控制流、递归、数据结构和宏等语言功能。然后,我们可以通过 LLVM 之类的编译器生成 SSA 形式的伴随代码,并将传统编译器优化的所有优势应用于前向和后向传播。此外,这种方法还为扩展该编译器基础结构提供了可能,可以使用更高级和特定领域的优化,例如用于 TPU 等加速器的内核融合和编译。TensorFlow 的 Swift 和 Myia 开发人员在源到源 AD 技术的复兴中正在探索类似的方法。

Julia 用于此任务的一个关键优势是它可用于实现基本数值计算库,如微分方程求解器或优化库;这巧妙地解决了机器学习社区不断增长的需求,研究人员通过高性能代码(如光线追踪和物理引擎)进行反向传播,但求梯度仍必须在 C++中手动实现。相比之下,由于 Julia 的实现是用 Julia 编写的,因此可以轻松对从 ODE 到金融定价模型等求微分。将这些强大的工具带入模型是深度学习真正成为可微分编程的关键。

编译 Julia 到 GPU 上

GPU 编程是现代机器学习的重要组成部分,但 GPU 通常被视为实现细节。因为框架在内部提供内核,但用户只能使用一组有限的数学运算,无法直接对 GPU 进行编程。相比之下,Julia 中的 GPU 编程一直是一流的 CUDA 内核(可以很好地编写并从脚本或 notebook 中运行)。如下简单的向量加法内核看起来类似于 CUDA C:

但是,Julia 的类型特化可以在 GPU 上实现一组强大的附加抽象。例如,上面的代码不限于浮点数的密集数组,而是可以给出复数的稀疏数组;Julia 的常规特化机制将动态地生成一组新的 PTX 指令。我们甚至可以将此代码进一步抽象为可利用「+」函数的「高阶内核」,从而在四行代码内创建一整套函数 map(f,x,y)。

这可以实现一些强大的技巧,即使你自己从不编写 CUDA 代码。例如,我们可以透明地将大型广播(broadcast)表达式(例如 1 /(1 + exp(-x))及其向后传递融合到单个 GPU 内核中,从而获得显着加速。我们期望原生 GPU 代码生成能力和生态系统将为各种基于 Julia 的机器学习库提供支持。

编译 Julia 到 TPU 上

更进一步,谷歌最近开放了云 TPU 使用的 XLA IR,使得其他框架和用户都可以利用这个重量级硬件。XLA 功能强大但有限制:它无法运行 Python 解释器,当然也没有良好的性能。

而我们只需要从编写的 Julia 程序中提取「静态图」并将其直接编译为 XLA,从而允许 Julia 本身在 TPU 上运行。(事实上,这只是 Julia 一般编译过程的简单扩展,它在将程序发送到 LLVM 之前从程序中提取最大的「静态子图」。)这使我们可以充分利用 Julia 语言的表现力,包括控制流、递归、多调度、高阶函数、强大的数据结构和抽象、自定义数值类型,以及现有的包,如微分方程求解器和线性代数例程。所有这些都在获得高性能收缩阵列引擎的优势的同时,在 TPU 内运行。你今天就可以尝试,其中包括 ResNet 等大型机器学习模型和 TSVD 等线性代数例程。

项目地址:https://github.com/JuliaTPU/XLA.jl

自动批处理(Automatic Batching)

为了从这些加速器中获得最大收益(每个内核启动可能会产生大量开销,但是在输入大小上可以很好地扩展),批处理程序通常会同时将前向和反向传播应用于多个训练样本。在简单的情况下,例如使用卷积网络,通过在额外的批量维度上拼接 10 张图像来处理这个问题会变得很简单。但是,当处理可变结构的输入(例如树或图形)时,此任务变得更加困难。

大多数研究人员通过人工完成批处理代码来解决这个问题,这样做的成本非常大。人们已经针对不同的框架提出了不同的解决方案(DyNet、TensorFlow Fold,它试图在可能的情况下将一些高级 OP 一起批处理,但是这些通常要么具有其自身的可用性问题,要么没有实现手写代码的性能。

我们认为这个问题与单程序多数据(SPMD)编程的问题完全相同,单程序多数据编程几十年来一直被语言和编译器社区充分研究。实际上,它与 GPU 内部使用的并行模型非常相似,并且已经实现 CPU 的 SIMD 单元的编译器变换。通过从这项工作中汲取灵感,我们在 Julia 中实现了相同的变换,为标量 SIMD 单元和模型级批处理提供 SPMD 编程。这使我们能够编写对单个样本进行操作的简单代码,同时仍然在现代硬件上获得最佳性能。

结论

我们相信机器学习的未来取决于编程语言和编译器技术,尤其是扩展新的或现有的语言以满足机器学习研究的高要求。这不仅适用于机器学习社区,也适用于一般的数值规划;能够支持微分、向量化和新型硬件的编程语言将足以推动科学的许多进步。

  • 发表于:
  • 原文链接:https://kuaibao.qq.com/s/20181209A0HL5G00?refer=cp_1026
  • 腾讯「云+社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。

扫码关注云+社区

领取腾讯云代金券