如何打开张量流xla

Posted 2023-03-06

【中文标题】如何打开张量流xla

【英文标题】：How to open tensorflow xla

【发布时间】：2018-10-19 10:00:17

【问题描述】：

我尝试了 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没有区别。

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

还有其他方法可以打开吗？或者它已经默认打开了，但是我用错了。

打开 xla 后是否有基准来分析加速？

谢谢~

【参考方案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，甚至 tf-nightly。如果这不起作用，是时候交叉手指和build from source。

您在代码中的哪个位置启用了 XLA？

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

例如：

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

【参考方案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