CODEWITHSUNDEEP
tensorflowkerasdeep-learninggoogle-colaboratorytpu
JP Zhang
JP Zhang

Reputation: 816

How to use TPU in TensorFlow custom training loop?

I don't use Keras. And I want to use TPUs on Google Colab. Questions:

  1. Can tf.session automatically use TPUs?
  2. What do tf.contrib.tpu.TPUDistributionStrategy, tf.contrib.tpu.rewrite, tf.contrib.cluster_resolver.TPUClusterResolver do in TPU computing? Are they all necessary?

Upvotes: 0

Views: 633

Answers (1)

Winnie Xu
Winnie Xu

Reputation: 31

Which Tensorflow version are you running? Currently the firmware for TPUs on Google Colab only support 1.14 (I may be wrong with the exact version, but it's definitely 1.x), however if you are using TF2.0, there is TPU support for nightly-2.x on GCP, so perhaps you can give that a try!

Note that in 2.0, you would want to get rid of any "sessions" because that is no longer a thing. Check out the TPUStrategy docs here for more information: https://www.tensorflow.org/guide/distributed_training#tpustrategy

Upvotes: 3

Related Questions