13

Google colab brings TPUs in the Runtime Accelerator. I found an example, How to use TPU in Official Tensorflow github. But the example not worked on google-colaboratory. It stuck on following line:

tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)

When I print available devices on colab it return [] for TPU accelerator. Does anyone knows how to use TPU on colab?

enter image description here

Bob Smith
  • 36,107
  • 11
  • 98
  • 91
Amir
  • 16,067
  • 10
  • 80
  • 119
  • 1
    I tried training a model with keras using the TPU runtime but it was really slow. What the GPU runtime trained in 2 minutes, the TPU runtime trained in 58 minutes. I'm waiting for a fix too. – Se7eN Sep 27 '18 at 15:38
  • 3
    To use the TPU, uou'll need to connect to it explicitly using the recipe in the answer below. (The use of TPUs isn't automatic like GPUs, unfortunately.) – Bob Smith Sep 27 '18 at 15:44

1 Answers1

18

Here's a Colab-specific TPU example: https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpu_and_keras.ipynb

The key lines are those to connect to the TPU itself:

# This address identifies the TPU we'll use when configuring TensorFlow.
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']

...

tpu_model = tf.contrib.tpu.keras_to_tpu_model(
training_model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
    tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

(Unlike a GPU, use of the TPU requires an explicit connection to the TPU worker. So, you'll need to tweak your training and inference definition in order to observe a speedup.)

Bob Smith
  • 36,107
  • 11
  • 98
  • 91
  • Thanks for your answer. I add **TPU_WORKER** line and resolved the running error. but I didnt gain any significant improvement. It's too slow. – Amir Sep 27 '18 at 17:01
  • 1
    My recommendation is to ask a separate question that includes a link to your notebook. Often, existing models require some adjustment in order to speed things up, and it's easy to end up in a case where critical operations are performed on the CPU rather than TPU. – Bob Smith Sep 27 '18 at 17:03
  • Thanks, It was a bug in my code and nothing to do with TPU. – Amir Sep 27 '18 at 17:27
  • This example only shows how to do it on keras. Are there any possible methods for "pure" tensorflow or pytorch? – Semin Park Sep 28 '18 at 01:14
  • I noticed even the network speed is faster with GPU environment compared to CPU and TPU. I noticed downloading dataset from Kaggle. – RomRoc Oct 10 '18 at 07:40
  • I am getting error as AttributeError: 'Activation' object has no attribute 'outbound_nodes' when i try to use TPU on google colab. – Gowtham M Oct 17 '18 at 00:30
  • Hi folks, please ask distinct questions for distinct issues. This question is about the basics of connecting to the TPU rather than performance optimization or error debugging. – Bob Smith Oct 17 '18 at 02:17