python - 如何在 Tensorflow 2.0 中使用 Google Colab 的 TPU？

Question

我正在尝试利用 Google Colab 使用张量处理单元 (TPU) 来训练神经网络。 Tensorflow 刚刚发布了一个主要版本 2.0，所以我尝试在 Tensorflow 2.0 中做到这一点。我尝试了以下三个指南，但所有指南都是为 Tensorflow 1.14 编写的，并且在 Tensorflow 2.0 中失败:

1) 遵循指南 TPUs in Colab ，我收到错误:

AttributeError: module 'tensorflow' has no attribute 'Session'

(来自引用:以 tf.Session(tpu_address) 作为 session :)

2) 遵循指南 Simple Classification Model using Keras on Colab TPU ，我得到同样的错误

3) 遵循指南 cloud_tpu_custom_training ，我收到错误:

AttributeError: module 'tensorflow' has no attribute 'contrib'

(来自引用:resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu=TPU_WORKER))

有人有使用 TPU 在 Tensorflow 2.0 中训练神经网络的示例吗？

编辑:这个问题似乎也在 github 上提出:InvalidArgumentError: Unable to find a context_id matching the specified one #1

Answer 1

首先，教程中给出的代码与 2.x 不兼容

1. 您需要在colab中选择runtime为TPU才能执行TPU中的代码

2. 对于错误

   AttributeError: module 'tensorflow' has no attribute 'Session'

   您需要使用tf.compat.v1.Session()，因为tf.session已弃用。

3. 请使用 tf.distribute.cluster_resolver 代替 tf.contrib.cluster_resolver

请引用Tensorflow Addon-repo将代码从 1.x 转换为 2.x 兼容。