如何在 TensorFlow 自定义训练循环中使用 TPU？

【问题标题】：How to use TPU in TensorFlow custom training loop?如何在 TensorFlow 自定义训练循环中使用 TPU？
【发布时间】：2020-02-22 03:15:28
【问题描述】：

我不使用 Keras。我想在 Google Colab 上使用 TPU。问题：

tf.session 可以自动使用 TPU 吗？
tf.contrib.tpu.TPUDistributionStrategy、tf.contrib.tpu.rewrite、tf.contrib.cluster_resolver.TPUClusterResolver 在 TPU 计算中有什么作用？它们都是必需的吗？

【问题讨论】：

【解决方案1】：

您正在运行哪个 TensorFlow 版本？目前 Google Colab 上的 TPU 固件仅支持 1.14（我可能对确切的版本有误，但肯定是 1.x），但是如果您使用的是 TF2.0，则在 GCP 上 TPU 支持 nightly-2.x ，所以也许你可以试试看！

请注意，在 2.0 中，您可能希望摆脱任何“会话”，因为这不再是一回事。在此处查看 TPUStrategy 文档以获取更多信息：https://www.tensorflow.org/guide/distributed_training#tpustrategy

【讨论】：