Julia语言在大数据和人工智能领域前景如何?
是否能超越Python不好说,不过Julia 语言发展确实非常迅速,它可以视为同时具备了 Python 的灵活性与 C 的速度,但目前 TensorFlow 和 PyTorch 等框架官方都不支持 Julia 语言。因此近日有研究者借助 XLA 底层编译器为 Julia 构建 TPU 支持,他们表示该方法能够将 Julia 程序编写的 VGG19 模型融合到 TPU 可执行文件中,并调用 TPU 实现高效计算。而 Google.ai 的负责人 Jeff Dean 在推特上也表示「Julia + TPUs = fast and easily expressible ML computations!」
1. 引言过去的几年里推动机器学习技术稳步发展的根本性改变之一是训练和优化机器学习模型的巨大计算力。许多技术都是很年前就已经提出,唯有近几年提升的计算力可以为现实世界的问题提供足够优质的解决方案。这些计算能力的很大一部分是通过 GPU 获取的,其针对向量的计算能力最初是为图形而设计的,但机器学习模型通常需要执行复杂的矩阵运算,因此 GPU 同样表现出了非常好的性能。
这些方法及 GPU 在现实世界,尤其是在机器学习领域的成功引发了硬件设计者的一系列创新,他们致力于为机器学习工作负载研发新的加速器。然而,尽管 GPU 很长一段时间都在 CUDA 等软件系统发力,但这些库通常不会扩展到新的非 GPU 加速器,为这些加速器开发软件仍然是一大挑战。
2017 年,谷歌宣布他们将通过云服务向大众提供他们专有的 TPU 机器学习加速器。最初,TPU 的使用局限于根据谷歌 TensorFlow 机器学习框架编写的应用。幸运的是,2018 年 9 月,谷歌通过底层 XLA(Accelerated Linear Algebra)编译器的 IR 开放了 TPU 的访问权限。这个 IR 是一个通用的优化编译器,用于表达线性代数基元的任意计算,因此为使用 TPU 的非 TensorFlow 用户以及非机器学习工作负载提供了良好的基础。
在本文中,我们介绍了使用这个接口编译通用 Julia 代码的初步工作,它们可以进一步访问谷歌云的 TPU。这一方法与 TensorFlow(Abadi et al., 2016)采用的方法形成对比,后者没有编译 python 代码,而是先用 Python 构建一个计算图,然后再对这个计算图进行编译。它在美学上类似于 JAX(Frostig et al., 2018),JAX 的目标是通过跟踪和 Offload 高级数组运算来 Offload Python 本身编写的计算。然而重要的是,我们不依赖于追踪,而是利用 Julia 的静态分析和编译能力来编译整个程序,包括传递到设备端的所有控制流。
值得一提的是,我们的方法允许用户在编写模型时充分利用 Julia 语言的表现力。这些表现力主要体现在一些高级特征上,如多重派发、高阶函数和现有库,如微分方程求解器(Rackauckas & Nie,2017)和通用线性代数例程等。由于只在纯 Julia 代码上运行,所以它也与 Zygote.jl(Innes, 2018)自动微分工具兼容,该工具能执行自动微分作为高级编译过程。总的来说,我们能够编译使用 Flux 机器学习框架编写的完整机器学习模型,将模型的前向、反向传播及训练回路融合成一个可执行文件,并 Offload 到 TPU 中。
论文:Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs
论文链接:https://arxiv.org/abs/1810.09868
摘要:谷歌的云 TPU 是一种前景广阔的新型机器学习工作负载硬件架构,近年来已经成就了谷歌很多里程碑式的机器学习突破。如今,谷歌已经在其云平台上为大众提供 TPU,最近又进一步开放,允许非 TensorFlow 前端使用。我们描述了一种通过这一新 API 及谷歌 XLA 编译器将 Julia 程序的适当部分 Offload 到 TPU 的方法和实现。我们的方法能够将 Julia 程序编写的 VGG19 模型及其正向传播完全融合到单个 TPU 可执行文件中,以便 Offload 到设备上。我们的方法与 Julia 代码上现有的基于编译器的自动微分技术很好地结合在一起,因此也能够自动获得 VGG19 反向传播并采用类似的方法将其 Offload 到 TPU。使用我们的编译器访问 TPU,我们能够在 0.23 秒内完成批量为 100 张图像的 VGG19 前向传播,而 CPU 上的原始模型则需要 52.4s。我们的实现仅需不到 1000 行的 Julia 代码,无需根据 TPU 对核心 Julia 编译器或任何其他 Julia 包进行特有的更改。
7 结果7.4 在 TPU 上进行评估
图 2:不同批大小对应的 VGG19 前向传播时长。
图 3:被编译为 XLA 后,Metalhead.jl VGG19 前向传播和反向传播的指令数分解。