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

Question

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

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

Answer 1

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

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