1.引言
过去的几年里推动机器学习技术稳步发展的根本性改变之一是训练和优化机器学习模型的巨大计算力。许多技术都是很年前就已经提出，唯有近几年提升的计算力可以为现实世界的问题提供足够优质的解决方案。这些计算能力的很大一部分是通过GPU获取的，其针对向量的计算能力最初是为图形而设计的，但机器学习模型通常需要执行复杂的矩阵运算，因此GPU同样表现出了非常好的性能。
这些方法及GPU在现实世界，尤其是在机器学习领域的成功引发了硬件设计者的一系列创新，他们致力于为机器学习工作负载研发新的加速器。然而，尽管GPU很长一段时间都在CUDA等软件系统发力，但这些库通常不会扩展到新的非GPU加速器，为这些加速器开发软件仍然是一大挑战。
2017年，谷歌宣布他们将通过云服务向大众提供他们专有的TPU机器学习加速器。最初，TPU的使用局限于根据谷歌TensorFlow机器学习框架编写的应用。幸运的是，2018年9月，谷歌通过底层XLA（AcceleratedLinearAlgebra）编译器的IR开放了TPU的访问权限。这个IR是一个通用的优化编译器，用于表达线性代数基元的任意计算，因此为使用TPU的非TensorFlow用户以及非机器学习工作负载提供了良好的基础。
在本文中，大家介绍了使用这个接口编译通用Julia代码的初步工作，它们可以进一步访问谷歌云的TPU。这一方法与TensorFlow（Abadietal.,2016）采用的方法形成对比，后者没有编译python代码，而是先用Python构建一个计算图，然后再对这个计算图进行编译。它在美学上类似于JAX（Frostigetal.,2018），JAX的目标是通过跟踪和Offload高级数组运算来OffloadPython本身编写的计算。然而重要的是，大家不依赖于追踪，而是利用Julia的静态分析和编译能力来编译整个程序，包括传递到设备端的所有控制流。
值得一提的是，大家的方法允许用户在编写模型时充分利用Julia语言的表现力。这些表现力主要体现在一些高级特征上，如多重派发、高阶函数和现有库，如微分方程求解器（Rackauckas&Nie，2017）和通用线性代数例程等。由于只在纯Julia代码上运行，所以它也与Zygote.jl（Innes,2018）自动微分工具兼容，该工具能执行自动微分作为高级编译过程。总的来说，大家能够编译使用Flux机器学习框架编写的完整机器学习模型，将模型的前向、反向传播及训练回路融合成一个可执行文件，并Offload到TPU中。
论文：AutomaticFullCompilationofJuliaProgramsandMLModelstoCloudTPUs
论文链接：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.jlVGG19前向传播和反向传播的指令数分解。