CODEWITHSUNDEEP
tensorflowgoogle-colaboratorytensorflow2.0tpu
Andrey
Andrey

Reputation: 6377

How to check if TPU is available in Google Colab?

I am trying to choose a distribution strategy based on availability of TPU.

My code is as follows:

import tensorflow as tf
if tf.config.list_physical_devices('tpu'):
  resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
  tf.config.experimental_connect_to_cluster(resolver)
  tf.tpu.experimental.initialize_tpu_system(resolver)
  print("All devices: ", tf.config.list_logical_devices('TPU'))
  strategy = tf.distribute.experimental.TPUStrategy(resolver)
else:  # use default strategy
  strategy = tf.distribute.get_strategy()

But it doesn't work.

How can I identify TPU ?

Upvotes: 3

Views: 3150

Answers (1)

Andrey
Andrey

Reputation: 6377

The following code works:

import tensorflow as tf
try:
  resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
  tf.config.experimental_connect_to_cluster(resolver)
  tf.tpu.experimental.initialize_tpu_system(resolver)
  print("All devices: ", tf.config.list_logical_devices('TPU'))
  strategy = tf.distribute.experimental.TPUStrategy(resolver)
except ValueError:
  strategy = tf.distribute.get_strategy()

Upvotes: 7

Related Questions