CODEWITHSUNDEEP
pythontensorflowgoogle-colaboratorygoogle-cloud-tpu
Jann
Jann

Reputation: 1829

Running cloud TPU profiler in Google Colab environment

I am running a Google Colab notebook and am trying to capture TPU profiling data for use in TensorBoard, however I can't get capture_tpu_profile to run in the background while running my TensorFlow code.

So far I tried to run the capture process in the background with:

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

and 

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

Upvotes: 1

Views: 1306

Answers (2)

michaelb
michaelb

Reputation: 252

One way to do this is to use the TPUProfilerHook

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

Which runs the profiler as a session hook.

Example here https://github.com/tensorflow/tpu/blob/5d838047af0163bdf7b97b9404648dc2961c4b63/models/official/resnet/resnet_main.py#L699

Upvotes: 0

Jann
Jann

Reputation: 1829

Turns out a way to do this is to start the process from python directly like this (I also had to modify the parameter from --tpu to --service_addr):

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

the check=True makes the command raise an Exception if it fails.

Upvotes: 5

Related Questions