选自julialang
作者：MikeInnes等
机器之心编译
参与：刘晓坤、思源
本文基于NeurIPSMLSys的一篇论文《FashionableModellingwithFlux》，探讨开发者们如何使用Julia语言从头开始思考机器学习工具，并提供对于现代机器学习工具所需改进的一些见解，涉及新的可微分编程工具Flux、求梯度、支持GPU和TPU、自动批处理。为什么Julia式的机器学习不需要计算图呢？因为Julia的机器学习语法就是计算图。
鉴于机器学习（ML）对编程语言、编译器和生态系统的众多需求，现在已经有很多有趣的发展。不仅TensorFlow和PyTorch等现有系统间的权衡得不到解决，而且这两个框架都包含不同的「静态图」和「eagerexecution」接口，但它们的形式已经比以前更加清晰。与此同时，机器学习模型基本上是可微分算法的思想（通常称为可微分编程）已经流行起来。
当前的机器学习框架遇到了阻碍，很多已有的新项目都完全移除了计算图，从而使可微分编程成为主流。例如，由Theano团队开发的Myia可以求微分并编译Python的一个子集为高性能GPU代码。SwiftforTensorFlow作为Swift语言的扩展，它可以将兼容的函数编译为TensorFlow计算图。最后，Flux生态系统为Julia编译器提供了一些机器学习专用的工具，包括：first-classgradients、即时CUDA核编译、自动批处理（automaticbatching）以及对新硬件（例如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提供代码编译，以及自动批处理。
求梯度
推动反向模式求微分的极限，我们将此视为语言层面的问题。求微分是一种符号转换，属于编译器的领域。现有框架通过追踪（实际上是一种部分评估或抽象解释）来实现这一目标。人们引入了一种新的张量类型，它记录了所执行的所有基本数学运算，生成一个计算图（或符号表达式），其中删除了宿主语言的控制流和数据结构。然而，这给出了一个艰难的权衡：我们要么接受解释器的开销（eagerexecution），要么固定用户的控制流并限制可以构建的模型种类（静态图）。
反之，如果「计算图」就是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中运行）。如下简单的向量加法内核看起来类似于CUDAC：
functionkernel_vadd(a,b,c)i=(blockIdx().x-1)*blockDim().x+threadIdx().xc[i]=a[i]+b[i]returnend
但是，Julia的类型特化可以在GPU上实现一组强大的附加抽象。例如，上面的代码不限于浮点数的密集数组，而是可以给出复数的稀疏数组；Julia的常规特化机制将动态地生成一组新的PTX指令。我们甚至可以将此代码进一步抽象为可利用「+」函数的「高阶内核」，从而在四行代码内创建一整套函数map（f，x，y）。
这可以实现一些强大的技巧，即使你自己从不编写CUDA代码。例如，我们可以透明地将大型广播（broadcast）表达式（例如1/（1+exp（-x））及其向后传递融合到单个GPU内核中，从而获得显着加速。我们期望原生GPU代码生成能力和生态系统将为各种基于Julia的机器学习库提供支持。
编译Julia到TPU上
更进一步，谷歌最近开放了云TPU使用的XLAIR，使得其他框架和用户都可以利用这个重量级硬件。XLA功能强大但有限制：它无法运行Python解释器，当然也没有良好的性能。
而我们只需要从编写的Julia程序中提取「静态图」并将其直接编译为XLA，从而允许Julia本身在TPU上运行。（事实上，这只是Julia一般编译过程的简单扩展，它在将程序发送到LLVM之前从程序中提取最大的「静态子图」。）这使我们可以充分利用Julia语言的表现力，包括控制流、递归、多调度、高阶函数、强大的数据结构和抽象、自定义数值类型，以及现有的包，如微分方程求解器和线性代数例程。所有这些都在获得高性能收缩阵列引擎的优势的同时，在TPU内运行。你今天就可以尝试，其中包括ResNet等大型机器学习模型和TSVD等线性代数例程。
项目