TensorFlow:XLA 不是 运行 "Invalid argument: No _XlaCompile" 错误

TensorFlow: XLA not running with "Invalid argument: No _XlaCompile" errors

我正在试用此处描述的 XLA 教程:https://www.tensorflow.org/performance/xla/jit

我来自 https://raw.githubusercontent.com/tensorflow/tensorflow/r1.1/tensorflow/examples/tutorials/mnist/mnist_softmax_xla.py 运行 mnist_softmax_xla.py,具有以下选项:

TF_CPP_MIN_VLOG_LEVEL=2 TF_XLA_FLAGS='--xla_generate_hlo_graph=.*' python mnist_softmax_xla.py

不幸的是,我在输出中遇到了一堆 "Custom creator error: Invalid argument: No _XlaCompile for Const" 错误(对于所有其他类型的操作也是如此)。此外,没有创建 hlo_graph_xx.dot 文件(正如教程所说的那样)。

我的 python 安装是 Ubuntu 16.04 LTS 上的 Anaconda 4.3.1 (Anaconda3-4.3.1-Linux-x86_64.sh)。

TensorFlow 是使用以下命令从源代码编译的 1.1.0 版:

$ echo "deb [arch=amd64] http://storage.googleapis.com/bazel-apt stable jdk1.8" | sudo tee /etc/apt/sources.list.d/bazel.list
$ curl https://bazel.build/bazel-release.pub.gpg | sudo apt-key add -
$ sudo apt-get update && sudo apt-get install bazel

$ git clone https://github.com/tensorflow/tensorflow
$ cd tensorflow/
$ git checkout v1.1.0
$ ./configure 

Please specify the location of python. [Default is /home/ubuntu/anaconda3/bin/python]: 
Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native]: 
Do you wish to use jemalloc as the malloc implementation? [Y/n] 
jemalloc enabled
Do you wish to build TensorFlow with Google Cloud Platform support? [y/N] 
No Google Cloud Platform support will be enabled for TensorFlow
Do you wish to build TensorFlow with Hadoop File System support? [y/N] 
No Hadoop File System support will be enabled for TensorFlow
Do you wish to build TensorFlow with the XLA just-in-time compiler (experimental)? [y/N] y
XLA JIT support will be enabled for TensorFlow
Found possible Python library paths:
  /home/ubuntu/anaconda3/lib/python3.6/site-packages
Please input the desired Python library path to use.  Default is [/home/ubuntu/anaconda3/lib/python3.6/site-packages]

Using python library path: /home/ubuntu/anaconda3/lib/python3.6/site-packages
Do you wish to build TensorFlow with OpenCL support? [y/N] 
No OpenCL support will be enabled for TensorFlow
Do you wish to build TensorFlow with CUDA support? [y/N] 
No CUDA support will be enabled for TensorFlow
Configuration finished
............
INFO: Starting clean (this may take a while). Consider using --expunge_async if the clean takes more than several minutes.
...........
INFO: All external dependencies fetched successfully.

$ bazel build --config=opt //tensorflow/tools/pip_package:build_pip_package
$ bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
$ pip install /tmp/tensorflow_pkg/tensorflow-1.1.0-cp36-cp36m-linux_x86_64.whl 

为什么 XLA 不适用于此设置?

如何使用 XLA 安装 TensorFlow?

已解决!

原因是 XLA 教程中的一个有点晦涩的注释:

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.

这不适用于 CPU:

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

这适用于 CPU(以及 GPU):

jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
with jit_scope():
    ...