CODEWITHSUNDEEP
tensorflowmachine-learningkerasgoogle-cloud-platform
Boppity Bop
Boppity Bop

Reputation: 10463

Use multiple TPUs in Keras with TF

How do I setup Keras model to use multiple TPUs for training? All the examples use a single TPU address

Upvotes: 1

Views: 725

Answers (2)

Boppity Bop
Boppity Bop

Reputation: 10463

Apparently correct way to use multiple TPUs is to pass list of addresses into the resolver:

TPU_ADDRESS1 = 'grpc://10.240.1.2:8470'
TPU_ADDRESS2 = 'grpc://10.240.2.2:8470'

tpu_model = tf.contrib.tpu.keras_to_tpu_model(model, 
        strategy=tf.contrib.tpu.TPUDistributionStrategy(
            tf.contrib.cluster_resolver.TPUClusterResolver(tpu = [TPU_ADDRESS1,TPU_ADDRESS2])))

However TF 1.13 does not support more than 1 TPU

Upvotes: 1

Dr Yuan Shenghai
Dr Yuan Shenghai

Reputation: 1915

In keras, this is not officially provided YET.

But I know people are working very hard to develop this function for Keras.

At bottom tensorflow layer, they do provide some experimental function. See

https://www.tensorflow.org/guide/distribute_strategy

https://www.tensorflow.org/guide/using_tpu

https://cloud.google.com/ml-engine/docs/tensorflow/using-tpus

"Experimental support for Cloud TPUs is currently available for Keras and Colab."

In your Tensorflow program, you should use TPUClusterResolver to connect with the TPU gRPC server running on the TPU VM. The TPUClusterResolver returns the IP address and port of the Cloud TPU.

Assigning ops to TPUs To make use of the TPUs on a machine, you must use the TensorFlow TPUEstimator API, which inherits from the high-level TensorFlow Estimator API.

It may not be exactly what you wanted e.g a local TPU cluster. But you can follow their way to get started

Personally, I never try with complicated multi gpu/tpu solution before. We only do simple single GPU training in school research. Here is what I can find to help you

So go join their beta and good luck!

Upvotes: 1

Related Questions