Запуск облачного профилировщика TPU в среде Google Colab

Question

Я использую записную книжку Google Colab и пытаюсь получить данные профилирования TPU для использования в TensorBoard, однако мне не удается заставить capture_tpu_profile работать в фоновом режиме при выполнении кода TensorFlow.

До сих пор я пытался запустить процесс захвата в фоновом режиме с помощью:

!capture_tpu_profile --logdir=gs://<my_logdir> --tpu=$COLAB_TPU_ADDR &

и

!bg capture_tpu_profile --logdir=gs://<my_logdir> --tpu=$COLAB_TPU_ADDR

Answer 1

Один из способов сделать это - использовать TPUProfilerHook

https://github.com/tensorflow/tpu/blob/master/models/common/tpu_profiler_hook.py

, который запускает профилировщик как перехватчик сеанса.

Пример здесь https://github.com/tensorflow/tpu/blob/5d838047af0163bdf7b97b9404648dc2961c4b63/models/official/resnet/resnet_main.py#L699

Answer 2

Оказывается, способ сделать это состоит в том, чтобы запустить процесс непосредственно из python, как это (я также должен был изменить параметр с --tpu до --service_addr):

import subprocess
subprocess.Popen(["capture_tpu_profile","--logdir=gs://<my_logdir>", "--service_addr={}".format(os.environ['COLAB_TPU_ADDR'])])

check=True заставляет команду вызывать Исключение в случае неудачи.