Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

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

本文基于 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:

function kernel_vadd(a, b, c)
    i = (blockIdx().x-1) * blockDim().x + threadIdx().x
    c[i] = a[i] + b[i]
    return
end

但是,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://julialang.org/blog/2018/12/ml-language-compiler

工程Julia编程语言静态图
5
相关数据
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

机器学习技术

机器学习是人工智能的一个分支,是一门多领域交叉学科,涉及概率论、统计学、逼近论、凸分析、计算复杂性理论等多门学科。机器学习理论主要是设计和分析一些让计算机可以自动“学习”的算法。因为学习算法中涉及了大量的统计学理论,机器学习与推断统计学联系尤为密切,也被称为统计学习理论。算法设计方面,机器学习理论关注可以实现的,行之有效的学习算法。

调度技术

调度在计算机中是分配工作所需资源的方法。资源可以指虚拟的计算资源,如线程、进程或数据流;也可以指硬件资源,如处理器、网络连接或扩展卡。 进行调度工作的程序叫做调度器。调度器通常的实现使得所有计算资源都处于忙碌状态,允许多位用户有效地同时共享系统资源,或达到指定的服务质量。 see planning for more details

Julia技术

Julia 是MIT设计的一个面向科学计算的高性能动态高级程序设计语言,项目大约于2009年中开始,2018年8月JuliaCon2018 发布会上发布Julia 1.0。据介绍,Julia 目前下载量已经达到了 200 万次,且 Julia 社区开发了超过 1900 多个扩展包。这些扩展包包含各种各样的数学库、数学运算工具和用于通用计算的库。除此之外,Julia 语言还可以轻松使用 Python、R、C/C++ 和 Java 中的库,这极大地扩展了 Julia 语言的使用范围。

规划技术

人工智能领域的「规划」通常是指智能体执行的任务/动作的自动规划和调度,其目的是进行资源的优化。常见的规划方法包括经典规划(Classical Planning)、分层任务网络(HTN)和 logistics 规划。

光线追踪技术

在计算机图形学中,光线跟踪是一种渲染技术,用于通过将光的路径跟踪为图像平面中的像素并模拟虚拟对象对光线的接收效果来生成图像。 该技术能够产生非常高的视觉真实感,通常高于典型扫描线渲染方法,但计算成本更高。

TensorFlow技术

TensorFlow是一个开源软件库,用于各种感知和语言理解任务的机器学习。目前被50个团队用于研究和生产许多Google商业产品,如语音识别、Gmail、Google 相册和搜索,其中许多产品曾使用过其前任软件DistBelief。

部分评估技术

在计算中,部分评估是针对几种不同类型的专业化程序优化的技术。 最直接的应用是产生比原件运行速度更快的新程序,同时保证以相同的方式运行。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

线性代数技术

线性代数是数学的一个分支,它的研究对象是向量,向量空间(或称线性空间),线性变换和有限维的线性方程组。向量空间是现代数学的一个重要课题;因而,线性代数被广泛地应用于抽象代数和泛函分析中;通过解析几何,线性代数得以被具体表示。线性代数的理论已被泛化为算子理论。由于科学研究中的非线性模型通常可以被近似为线性模型,使得线性代数被广泛地应用于自然科学和社会科学中。

批次技术

模型训练的一次迭代(即一次梯度更新)中使用的样本集。

Jupyter技术

Jupyter Notebook(此前被称为 IPython notebook)是一个交互式笔记本,支持运行 40 多种编程语言。 Jupyter Notebook 的本质是一个 Web 应用程序,便于创建和共享文学化程序文档,支持实时代码,数学方程,可视化和 markdown。 用途包括:数据清理和转换,数值模拟,统计建模,机器学习等等 。

推荐文章
暂无评论
暂无评论~