如何打开tensorflow xla

Question

我在tensorflow / examples / tutorials / mnist下尝试了该示例并尝试应用xla进行加速。 但是，我看不到https://www.tensorflow.org/performance/xla/jit所说的XlaLaunch。

此外，我尝试通过以下方式分析执行时间：

  train_loops = 100000
  t_start = time.time()
  for i in range(train_loops):
    batch_xs, batch_ys = mnist.train.next_batch(100)

    # Create a timeline for the last loop and export to json to view with
    # chrome://tracing/.
    if i == train_loops - 1:
      sess.run(train_step,
               feed_dict={x: batch_xs,
                          y_: batch_ys},
               options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE),
               run_metadata=run_metadata)
      trace = timeline.Timeline(step_stats=run_metadata.step_stats)
      with open('timeline.ctf.json', 'w') as trace_file:
        trace_file.write(trace.generate_chrome_trace_format())
    else:
      sess.run(train_step, feed_dict={x: batch_xs, y_: batch_ys})

  tdiff = time.time() - t_start
  print("tdiff", tdiff, " i = ", i)

使用xla和不使用xla似乎没有什么区别。

我看到一些文章说我应该“重建” tensorflow源来打开xla？我应该吗？

还有其他打开方式吗？或默认情况下已将其打开，但使用它的方式有误。

打开xla后，是否有基准可用来分析加速情况？

谢谢〜

Answer 1

pip中可用的最新TensorFlow二进制文件可能包括已内置的XLA支持。 pip install tensorflow --upgrade --force-reinstall或pip install tensorflow-gpu --upgrade --force-reinstall将为您提供此支持（自TF 1.12起）。这对我适用于具有或不具有NVIDIA GPU支持的Ubuntu，但不适用于我的MacBook。如果您的二进制文件没有，则在尝试使用基准测试时会在下面知道。

如果您正确地调用XLA JIT，并且二进制文件中没有内置的支持，则会出现异常。

ValueError: Op type not registered 'XlaClusterOutput' in binary running on localhost. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) tf.contrib.resampler should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed. while building NodeDef 'tower_0/v/output0'

如果得到这个，请尝试安装最新的tensorflow，甚至每晚安装一次。如果那不起作用，请花些时间和build from source交叉手指。

您在代码中的什么地方启用了XLA？

对于测试基准，您可以尝试使用官方TensorFlow Benchmarks存储库here。如果未下载ImageNet，则可以使用合成数据。您可以使用标志--xla_compile=True或False启用或禁用XLA。

例如：

python tf_cnn_benchmarks.py --num_batches=2000 --xla_compile=True

Answer 2

无需从带有任何标志的源代码安装 Tensorflow 即可启用 XLA 的一种方法是遵循 https://github.com/tensorflow/tensorflow/issues/44683 中的说明，并在运行运行训练的脚本时设置环境变量：

>>> TF_XLA_FLAGS=--tf_xla_enable_xla_devices python my_training_script.py

然后您应该会看到日志消息：

XLA service 0x95522ao initialized for platform Host (this does not guarantee that XLA will be used). Devices:
  StreamExecutor device (0): Host, Default Version