CODEWITHSUNDEEP
deep-learninggoogle-colaboratorygoogle-cloud-tpu
Nick Ben
Nick Ben

Reputation: 63

Google Colab KeyError: 'COLAB_TPU_ADDR'

I'm trying to run a simple MNIST classifier on Google Colab using the TPU option. After creating the model using Keras, I am trying to convert it into TPU by:

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())

And the error I'm getting is:

---------------------------------------------------------------------------
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'

It looks like I need to change the TPU address but been googling and haven't found anything yet. Appreciate some help, thanks!

Upvotes: 3

Views: 10124

Answers (3)

Erwing Forero
Erwing Forero

Reputation: 41

Use the next code:

import tensorflow as tf
print("Tensorflow version " + tf.__version__)

try:
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
except ValueError:
  raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')

Upvotes: 0

Ash
Ash

Reputation: 1

Many times in Google Colab TPU is not able, and you won't get assigned a TPU. And Colab will throw this error. You will need to find a time of the day when its more likely to get Colab TPU, eg early morning or late at night.

Upvotes: 0

Bob Smith
Bob Smith

Reputation: 38579

You'll need to change the backend to include a TPU using the notebook settings available in the Edit -> Notebook settings menu.

enter image description here

Upvotes: 12

Related Questions