Train object detection model on Google Colab with TPU - tensorflow

Question

I'm training an object detection model by following the guide here https://towardsdatascience.com/creating-your-own-object-detector-ad69dda69c85
On Google Colab I am able to execute the following and it makes use of the GPU
python train.py --train_dir=training/ --pipeline_config_path=training/ssd_mobilenet_v1_0.75_depth_quantized_300x300_coco14_sync.config
I would now like train by using the TPU but this obviously does not just work out of the box. Running train.py is slow and appears to be using CPU only. How can I achieve this?

Answer 1

While using TPU in Google Colab, we should use the below mentioned code tocheck that the TPU devices are properly recognized in the environment:
import os
import pprint
import tensorflow as tf
if 'COLAB_TPU_ADDR' not in os.environ:
print('ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!')
else:
tpu_address = 'grpc://' + os.environ['COLAB_TPU_ADDR']
print ('TPU address is', tpu_address)
with tf.Session(tpu_address) as session:
devices = session.list_devices()
print('TPU devices:')
pprint.pprint(devices)
This should output a list of 8 TPU devices available in our Colab environment.
In order to run the tf.keras model on a TPU, we have to convert it to a TPU-model using the tf.contrib.tpu.keras_to_tpu module.
It can be done using the below code:
# This address identifies the TPU we'll use when configuring TensorFlow.
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tf.logging.set_verbosity(tf.logging.INFO)
resnet_model = tf.contrib.tpu.keras_to_tpu_model(
resnet_model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))
Fore more information, please refer this Medium Link and this Link.