tensorflow - 使用 TensorFlow 2.1 的 Cloud TPU v3 Pod 的 TPUClusterResolver 出现错误

Question

我尝试在具有 TensorFlow 2.1 的 Google Cloud Compute Engine 虚拟机上使用我的(抢占式)Cloud TPU v3-256，但它似乎无法像 TPUClusterResolver 一样工作抛出无法查找 TPU 元数据错误。

只要我使用 grpc:// 地址而不是 TPU 名称，使用单独的(不可抢占的)TPU 就可以正常工作。但是，在使用 TPU 名称时，单个 TPU 和我的 TPU Pod 都不起作用，并抛出此错误。

有人可以帮我解决这个问题吗？

代码:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='my-tpu-name', zone='europe-west4-a', project='my-project')  # The zone, project and TPU Name are correct

输出:

ValueError: Could not lookup TPU metadata from name 'my-tpu-name'. Please double
check the tpu argument in the TPUClusterResolver constructor.
Exception: Failed to retrieve http://metadata.google.internal/computeMetadata/v1/
instance/service-accounts/default/?recursive=True
from the Google Compute Enginemetadata service. Response: {'metadata-flavor': 'Google', 
'date': 'Thu, 28 May 2020 17:42:35 GMT', 'content-type': 'text/html; charset=UTF-8',
'server': 'Metadata Server for VM', 'content-length': '1629', 'x-xss-protection': '0', 'x
frame-options': 'SAMEORIGIN', 'status': '404'}

Answer 1

我怀疑可能是以下任一方面不匹配:计算 VM 和 TPU 之间的 Tensorflow 版本、区域或项目。 如果您使用相同的 Tensorflow 版本(2.1 或 2.2)创建 TPU 和 GCE VM，并且它们都创建在同一项目和区域中。您只需在 TPUClusterResolver 中提供 TPU 名称即可，它应该可以正常工作:

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='my-tpu-name') 

如果您在虚拟机上设置 TPU_NAME 环境变量 (export TPU_NAME=my-tpu-name)，则可以省略 TPU 名称。