Как использовать ТПУ в Google Colab

Question

Google colab добавляет TPU в Runtime Accelerator.Я нашел пример, как использовать TPU в Официальном Tensorflow github .Но пример не сработал на гугл-колаборатории.Он застрял в следующей строке:

tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)

Когда я печатаю доступные устройства на colab, возвращается [] для ускорителя TPU.Кто-нибудь знает, как использовать ТПУ на колабе?

Answer 1

Вот пример конкретного TPU для Colab: https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpu_and_keras.ipynb

Основные линии для подключения к самому TPU:

# This address identifies the TPU we'll use when configuring TensorFlow.
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']

...

tpu_model = tf.contrib.tpu.keras_to_tpu_model(
training_model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
    tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

(В отличие от графического процессора, использование TPUтребуется явное соединение с работником TPU, поэтому вам нужно настроить свое обучение и определение логического вывода, чтобы наблюдать ускорение.)