tensorflow - 如何在谷歌colab中使用TPU

Question

Google colab 在运行时加速器中引入了 TPU。我找到了一个例子，How to use TPU in Official Tensorflow github。但该示例不适用于 google-colaboratory。它停留在以下行：

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

当我在 colab 上打印可用设备[]时，它会返回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。因此，您需要调整训练和推理定义以观察加速。）