如何在google colab中使用TPU

Question

Google colab在运行时加速器中引入了TPU。我找到了一个例子，如何在Official Tensorflow github中使用TPU。但是这个例子在google-colaboratory上不起作用。它停留在下面这一行：

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

当我在colab上运行print available devices时，它返回用于TPU加速器的[]。有人知道如何在colab上使用TPU吗？

​

​

Answer 1

下面是一个特定于Colab的TPU示例：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)))

(与GPU不同，使用TPU需要显式连接到TPU worker。因此，您需要调整您的训练和推理定义，以观察加速效果。)