Как открыть тензор потока xla

Question

Я попробовал пример в tenorflow / examples / tutorials / mnist и попытался применить xla для ускорения. Однако я не вижу XlaLaunch, как сказал https://www.tensorflow.org/performance/xla/jit.

Кроме того, я пытаюсь профилировать время выполнения с помощью:

  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.

Я видел некоторые статьи, в которых говорилось, что я должен "перестроить" исходный код tenorflow, чтобы открыть xla? нужно ли мне?

Есть ли другой способ открыть его? Или он был открыт по умолчанию, но я неправильно его использую.

Есть ли тест для профилирования ускорения после открытия xla?

Благодаря ~

Answer 1

Последние двоичные файлы TensorFlow, доступные в pip, могут включать встроенную поддержку XLA. pip install tensorflow --upgrade --force-reinstall или же pip install tensorflow-gpu --upgrade --force-reinstall окажет вам эту поддержку (начиная с версии 1.12). Это работает для меня в Ubuntu с поддержкой NVIDIA GPU или без нее, но не на моем MacBook. Если ваш двоичный файл этого не делает, вы узнаете ниже при попытке использовать тест.

Если вы правильно вызываете JIT XLA, и у вас нет встроенной поддержки двоичного кода, вы получите исключение. Что-то вроде

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'

Если вы получили это, попробуйте установить последний тензор потока или даже tf-nightly. Если это не сработает, пора скрестить пальцы и строить из источника.

Где в вашем коде вы включили XLA?

Для теста производительности вы можете попробовать воспользоваться официальным репозиторием TensorFlow Benchmarks здесь. Вы можете использовать синтетические данные, если у вас не загружен ImageNet. Вы можете включить или отключить XLA с флагом --xla_compile=True или же False,

Например:

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

Answer 2

Один из способов включить XLA без установки Tensorflow из источника с любыми флагами - следовать инструкциям в 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