如何在谷歌colab中使用TPU

Posted 2023-02-23

【中文标题】如何在谷歌colab中使用TPU

【英文标题】：How use TPU in google colab

【发布时间】：2019-03-02 23:45:56

【问题描述】：

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 吗？

【问题讨论】：

我尝试使用 TPU 运行时使用 keras 训练模型，但速度真的很慢。 GPU 运行时训练了 2 分钟，TPU 运行时训练了 58 分钟。我也在等修复。

要使用 TPU，您需要使用以下答案中的配方明确连接到它。 （不幸的是，TPU 的使用不像 GPU 那样自动。）

【参考方案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 工作器。因此，您需要调整训练和推理定义以观察加速。）

【讨论】：

感谢您的回答。我添加了 TPU_WORKER 行并解决了运行错误。但我没有得到任何显着的改善。太慢了。

我的建议是提出一个单独的问题，其中包含指向您笔记本的链接。通常，现有模型需要进行一些调整以加快速度，并且很容易导致关键操作在 CPU 而不是 TPU 上执行。

谢谢，这是我的代码中的一个错误，与 TPU 无关。

这个例子只展示了如何在 keras 上做。 “纯” tensorflow 或 pytorch 是否有任何可能的方法？

我注意到，与 CPU 和 TPU 相比，GPU 环境下的网络速度甚至更快。我注意到从 Kaggle 下载数据集。