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

Posted 2021-04-29 机器之心

选自julialang

机器之心编译

参与：刘晓坤、思源

本文基于 NeurIPS MLSys 的一篇论文《Fashionable Modelling with Flux》，探讨开发者们如何使用  语言从头开始思考机器学习工具，并提供对于现代机器学习工具所需改进的一些见解，涉及新的可微分编程工具 Flux、求梯度、支持 GPU 和 TPU、自动批处理。为什么 Julia 式的机器学习不需要计算图呢？因为 Julia 的机器学习语法就是计算图。

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

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

自动批处理（Automatic Batching）

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

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

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

结论

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

原文链接：https://julialang.org/blog/2018/12/ml-language-compiler