Google Colab KeyError：“COLAB_TPU_ADDR”

Question

我正在尝试使用 TPU 选项在 Google Colab 上运行一个简单的 MNIST 分类器。使用 Keras 创建模型后，我尝试通过以下方式将其转换为 TPU：

import tensorflow as tf
import os

tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
    )
)
tpu_model.compile(
    optimizer='rmsprop',
    loss='categorical_crossentropy',
    metrics=['accuracy']
)


print(model.summary())

我得到的错误是：

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-5-63c528142aab> in <module>()
      5     model,
      6     strategy=tf.contrib.tpu.TPUDistributionStrategy(
----> 7         tf.contrib.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
      8     )
      9 )

/usr/lib/python3.6/os.py in __getitem__(self, key)
    667         except KeyError:
    668             # raise KeyError with the original key value
--> 669             raise KeyError(key) from None
    670         return self.decodevalue(value)
    671 

KeyError: 'COLAB_TPU_ADDR'

看起来我需要更改 TPU 地址，但我在谷歌上搜索并没有找到任何东西。感谢您的帮助，谢谢！

Answer 1

您需要使用编辑 -> 笔记本设置菜单中的笔记本设置更改后端以包含 TPU。