TensorFlow XLA 学习笔记（一）：源码编译 + 运行例程 - 文章 - 开发者社区

picture.image

在构建和部署机器学习系统时，性能是至关重要的。因此，TensorFlow 中包含了 XLA，这是一款强大的线性代数编译器，可以帮助 TensorFlow 代码在嵌入式处理器、CPU、GPU、TPU 和其他硬件平台上尽可能快速地运行。

在前面文章《Nvidia Jetson TX2 上编译安装 TensorFlow r1.5》中介绍过源码编译 TensorFlow r1.5 的步骤，其中在 configure 阶段有一个选项：


        
Do you wish to build TensorFlow with XLA JIT support? [y/N]: y
XLA JIT support will be enabled for TensorFlow.

按上面选项配置完毕, bazel build 生成的 tensorflow 就包含了 XLA JIT 功能，可以继续下面试验。

在 TensorFlow 源码 tensorflow/examples/tutorials/mnist/ 目录下可找到 mnist_softmax_xla.py 这个例程。命令行选项如下：


        
usage: mnist_softmax_xla.py [-h] [--data_dir DATA_DIR] [--xla XLA]

optional arguments:
  -h, --help           show this help message and exit
  --data_dir DATA_DIR  Directory for storing input data
  --xla XLA            Turn xla via JIT on

运行该例程后会生成 timeline.ctf.json 文件，这个文件是由 TensorFlow 生成的 profiling 结果，可以通过 Chrome 浏览器打开查看。

Chrome 地址栏输入：chrome://tracing/，界面如下：

picture.image

点击“Load”，选择上面生成的 timeline.ctf.json 文件，界面如下：

picture.image

上图为关闭 XLA 时的结果。通过设置命令行选项打开 XLA，相应的界面如下：

picture.image

仔细对比可发现，关闭 XLA 时程序启动的 Kernel 名字为 MatMul、Add、SparseSoftmaxCrossEntropyWithLogits 等 TensorFlow 内置 OP Kernel 名，而打开 XLA JIT 时则改为 _XlaLaunch，这说明 XLA JIT 具有动态生成 Kernel 的能力，可以根据运行时实际硬件自动调优，加快运行速度。

如果希望在自己的代码中使用 XLA JIT 功能，请继续向下看。

在创建 Session 时，增加 config 参数，设置 config.graph_options.optimizer_options.global_jit_level 值为 tf.OptimizerOptions.ON_1 即可打开 XLA JIT 功能。 注：该配置对整个 Session 生效，所有 OP 都会受到影响。


        
  config = tf.ConfigProto()
  config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
  sess = tf.Session(config=config)

但上述配置仅适用于 GPU （使能 CUDA）的环境。在 CPU ONLY 模式下（参考【1】）

Note: Turning on JIT at the session level will not result in operations being compiled for the CPU. JIT compilation for CPU operations must be done via the manual method documented below. This decision was made due to the CPU backend being single-threaded.

下面这种开启 XLA JIT 的方式 可选择性地针对某几个 OP 生效 。


        
    jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
    x = tf.placeholder(np.float32)
    with jit_scope():
      y = tf.add(x, x)  # The "add" will be compiled with XLA.

只有位于 jit_scope with 块内的 OP 会选择 XLA JIT 优化，其他 OP 不受影响。这有利于精细控制优化范围，避免负优化。

第三种开启 XLA JIT 的方式：将需要优化的 OP 放置到 XLA Device（XLA_CPU 或 XLA_GPU）上。


        
with tf.device("/job:localhost/replica:0/task:0/device:XLA_GPU:0"):
  output = tf.add(input1, input2)

注意第三种方式可能会有额外的数据拷贝开销。

下面我们跑个 AlexNet Benchmark 【2】看看 XLA JIT 的威力如何。

硬件环境：GTX 1080

软件环境：TensorFlow 1.5 + CUDA 9.0 + cuDNN 7

picture.image

可见在使能 XLA JIT 时，Forward 速度提升 3X！Forward + Backward 速度提升 2.6X！说明 XLA JIT 对于训练和预测都有极大性能提升。

后面文章会深入研究 XLA 特性。

参考资料

【1】 https://www.tensorflow.org/performance/xla/jit

【2】 https://github.com/tensorflow/models/blob/master/tutorials/image/alexnet/alexnet\_benchmark.py

picture.image