β

XLA – TensorFlow 编译器

谷歌中国开发者社区 (GDG) 191 阅读

code { background-color: transparent }由 Google 内部 XLA 团队与 TensorFlow 团队合作发布

TensorFlow 的其中一个设计目标和核心优势是它的灵活性。TensorFlow 的设计目标是成为一个灵活并且可扩展的系统,用于定义任意数据流图,以及利用异构计算设备(例如 CPU 和 GPU)分布式地高效执行它们。


但灵活性往往与性能存在矛盾。尽管 TensorFlow 旨在让您定义任何种类的数据流图,但让所有图表都高效执行颇具挑战,因为 TensorFlow 会单独优化每一个运算。如果运算具有高效实现,或者每个运算都是相对重量级的运算,一切还好;否则,虽然用户仍可利用低级别运算组合该运算,却不能保证这个组合出的运算能够以最高效的方式运行。


正因如此,我们才会开发 XLA (加速线性代数)这个 TensorFlow 编译器。XLA 利用 JIT 编译技术分析用户在运行时创建的 TensorFlow 图表,根据实际运行时维度和类型将其专门化,将多个运算融合在一起并为它们生成高效的本机代码——适用于 CPU、GPU 之类的设备和自定义加速器(例如,Google 的 TPU)。

融合可组合运算来提升性能

我们以 tf.nn.softmax 运算为例。它计算其参数 Softmax 激活的公式如下:

CodeCogsEqn.gif


Softmax 的实现形式可以是 TensorFlow 基元运算(指数运算、归约运算、对应元素除法运算等)的组合运算:

softmax = exp(logits) / reduce_sum(exp(logits), dim)


由于需要额外移动数据,以及需要对运算之外不需要的临时结果进行实体化处理,这种运算的效率可能较低。不仅如此,在 GPU 之类的协处理器上,此类分解实现可能导致多次“内核启动”,令运算效率进一步下降。


XLA 好比秘制编译器调味酱,可以帮助 TensorFlow 自动优化基元运算的组合。在 XLA 的加持下,Tensorflow 能够在运行时分析图表,将运算融合在一起,并为融合的子图表生成高效的机器码,从而既保留灵活性,又不会牺牲运行时性能。


例如,上面介绍的 Softmax 分解实现经过 XLA 优化后,效率可提升至与手动优化复合运算相同的水平。


更概括地讲,XLA 可将 TensorFlow 运算的整个子图融合成内核启动次数最少的高效循环。例如:


上图中的许多运算都可融合成单个元素对应循环。以将 matmul 结果中的单个元素与 bias 矢量单个元素相加为例。此加法运算的结果是可与 0 进行比较的单个元素(用于 ReLU)。比较的结果可以进行乘方运算并与所有输入的指数之和进行除法运算,得到 Softmax 的输出。我们实际上并不需要在内存中创建 matmul、add 和 ReLU 的中间数组。

s[j] = softmax[j](ReLU(bias[j] + matmul_result[j]))


融合实现无需分配多余的内存就能在单个元素对应循环内计算出最终结果。在更高级的情境下,甚至可以将这些运算融合成矩阵乘法运算。

XLA 帮助 TensorFlow 在保留其灵活性的同时不必担心性能。


内部基准测试数据表明,在 Nvidia GPU 上,与不使用 XLA 的 TensorFlow 相比,使用 XLA 的 TensorFlow 速度提升高达 50%。不出所料,加速效果最好的那些模型都具有可融合成高效循环的长序列元素对应运算。不过,仍应将 XLA 视为实验性技术,在某些基准测试中可能会遭遇效率下降的情况。

https://www.youtube.com/watch?v=kAOanJczHA0
在上面这段 TensorFlow 开发者峰会演讲中,Chris Leary 和 Todd Wang 介绍了 TensorFlow 如何利用 XLA、JIT、AOT 以及其他技术最大限度缩短执行时间和增加计算资源。


通过极端专门化减小可执行文件大小

除了性能提升外,由于 XLA 能够减小可执行文件大小,TensorFlow 模型还能在内存受限环境中(例如移动设备上)受益。 tfcompile 是一个充分利用 XLA 进行提前 (AOT) 编译的工具:一个完整图表编译到 XLA,然后由后者生成紧凑机器码,在图表中实现运算。与最短运行时相结合,此架构可显著减小可执行文件大小。

例如,假定有一个基于 Android ARM 架构、深 3、宽 60 的堆叠 LSTM 模型,原始 TF 模型大小为 2.6 MB(1 MB 运行时 + 1.6 MB 图表);使用 XLA 编译后,大小减少到 600 KB。

点击此处可查看分辨率更高的图片。

如此大幅度的大小缩减是通过对模型进行其静态编译所暗示的完全专门化实现的。当模型运行时,不需要全部发挥 TensorFlow 运行时的能力和灵活性——只有实现用户真正感兴趣的图表的运算才会编译到原生代码中。尽管如此,XLA CPU 后端生成的代码性能仍远未达到最佳水平;项目的这个部分需要进一步优化。


对备选后端和设备的支持


为了让 TensorFlow 图表能够在现今的新型计算设备上执行,需要为新设备重新实现所有 TensorFlow 运算(内核)。视设备而定,这项工作可能非常繁重。

XLA 通过在设计上增加自定义后端,大大简化了对新设备的支持。由于 TensorFlow 可以 XLA 为目标,因此为 XLA 添加新设备后端后,XLA 便可运行 TensorFlow 图表。XLA 为新设备提供的实现表面小得多,因为 XLA 运算 不过是基元运算(还记得吧,XLA 独立处理复杂运算的分解)。我们 在此页面上 记录了为 XLA 添加自定义后端的过程。Google 利用此机制从 XLA 锁定 TPU

结论和展望

XLA 仍处于早期开发阶段。它在某些用例下呈现的结果很有前景,TensorFlow 未来可进一步受益于这项技术,这一点毋庸置疑。我们之所以决定提前将 XLA 发布到 TensorFlow Github ,是为了向社区征集代码,以及提供一个方便的表面,以便针对各种计算设备对 TensorFlow 进行优化,以及重构 TensorFlow 运行时和模型,让它们能够在新型硬件上运行。

<!—->

Source: XLA – TensorFlow 编译器

作者:谷歌中国开发者社区 (GDG)
Google Developer Groups
原文地址:XLA – TensorFlow 编译器, 感谢原作者分享。

发表评论