首页 > 编程笔记

TensorFlow XLA加速线性代数编译器

加速线性代数器(Accelerated linear algebra,XLA)是线性代数领域的专用编译器。根据 https://www.tensorflow.org/performance/xla/,它仍处于实验阶段,用于优化 TensorFlow 计算。

XLA 可以提高服务器和移动平台的执行速度、内存使用率和可移植性。它提供了双向 JIT(Just In Time)编译或 AoT(Ahead of Time)编译。使用 XLA,你可以生成平台相关的二进制文件(针对大量平台,如 x64、ARM等),可以针对内存和速度进行优化。

准备工作

目前,XLA 并不包含在 TensorFlow 的二进制版本中。用时需要从源代码构建它。

从源代码构建 TensorFlow,需要 TensorFlow 版的 LLVM 和 Bazel。TensorFlow.org 仅支持从 macOS 和 Ubuntu 的源代码构建。从源代码构建 TensorFlow 所需的步骤如下(参见https://www.tensorflow.org/install/install_sources):
  1. 确定要安装哪个版本的 TensorFlow——仅支持 CPU 的 TensorFlow 或支持 GPU 的 TensorFlow。
  2. 复制 TensorFlow 存储库:


     
  3. 安装以下依赖:
    • Bazel
    • TensorFlow 的 Python 依赖项
    • 对GPU版本,需要NVIDIA软件包以支持TensorFlow
  4. 配置安装。在这一步中,需要选择不同的选项,如 XLA、Cuda 支持、Verbs 等:

    ./configure

  5. 使用 bazel-build。
  6. 对于仅使用 CPU 的版本:


     
  7. 如果有兼容的 GPU 设备,并且需要 GPU 支持,请使用:


     
  8. 成功运行后,将获得一个脚本:build_pip_package。按如下所示运行这个脚本来构建 whl 文件:


     
  9. 安装 pip 包:


现在你已经准备好了。

具体做法

TensorFlow 生成 TensorFlow 图表。在 XLA 的帮助下,可以在任何新类型的设备上运行 TensorFlow 图表。
要进行相同的高级步骤,可以参考https://www.tensorflow.org/performance/xla/tfcompile

关注公众号「站长严长生」,在手机上阅读所有教程,随时随地都能学习。本公众号由站长亲自运营,长期更新,坚持原创,专注于分享创业故事+学习历程+工作记录+生活日常+编程资料。

公众号二维码
微信扫码关注公众号

优秀文章