Как проверить, доступен ли TPU в Google Colab?

Я пытаюсь выбрать стратегию распространения в зависимости от наличия TPU.

Мой код выглядит следующим образом:

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() 

Но не работает.

Как я могу идентифицировать ТПУ?


person Andrey    schedule 03.07.2020    source источник


Ответы (1)


Следующий код работает:

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() 
person Andrey    schedule 04.07.2020