How to convert tf.Keras model to TPU using TensorFlow 2.0 in Google Colab?

Question

Since TF 2.0 doesn't have the tf.contrib layer, how do I convert my model to run the training on TPU, without having access to tf.contrib.tpu.keras_to_tpu_model()?

I tried looking for code, but all of it are running on TensorFlow 1.x

My data is in .npy, I have a simple model, and I'm using only model.compile() and model.fit() to train it, but it looks like the model is running on CPU (taking 30 minutes/epoch vs the 2 minutes/epoch on GPU).

Answer 1

Colab with TensorFlow 2.2 (Updated Mar 2020)

It works after I fixed this issue, there's also a Colab Notebook at here

Convert Keras Model to TPU with TensorFlow 2.0 (Update Nov 2019)

To use Keras Model with Google Cloud TPU is very easy with TensorFlow 2.0, it does not need to be "converted" anymore.

What we need to do is just specify a Distributed Strategy to make the TensorFlow do all the heavy lift for us.

def create_model():
    model = tf.keras.models.Sequential()

    model.add(tf.keras.layers.Conv2D(128, (3, 3), input_shape=x_train.shape[1:]))
    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
    model.add(tf.keras.layers.Activation('elu'))

    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(10))
    model.add(tf.keras.layers.Activation('softmax'))

    return model

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_host(resolver.master())
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

with strategy.scope():
    model = create_model()
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
        loss=tf.keras.losses.sparse_categorical_crossentropy,
        metrics=[tf.keras.metrics.sparse_categorical_accuracy])

We create a Distribute Resolver, then create the Strategy for the Resolver, then with the strategy.scope(), we create our Keras Model, then we are done.

Learn more about how to create a Keras Model with TPU from my Colab notebook at https://colab.research.google.com/github/huan/tensorflow-handbook-tpu/blob/master/tensorflow-handbook-tpu-example.ipynb.

The Official TensorFlow Distributed Training document: https://www.tensorflow.org/guide/distributed_training

However, please notice that there are some environmental issues with the Colab that need to be fixed, so it might not be able to run in the Colab yet.

Colab TPU is not ready with 2.0 (yet)

The TensorFlow 2.0 had been released, but the Colab with TPU support still does not be supported yet. It was said by some Googlers that it will be ok after the 2.1 released.

We have an issue for tracking this progress at; https://github.com/huan/tensorflow-handbook-tpu/issues/1#issuecomment-552807573

My Old Answer

It was confirmed from Googler Wolff that we can not use TF 2.0 in Colab with TPU yet(reported at 15th April 2019):

The TPU you will get allocated via Colab is running TF 1.x. When you install the nightly 2.0 pip on your Jupyter VM, it doesn't change the TPU. You end up with a mismatch between what's running on your Jupyter instance and what the TPU has.

And according to https://github.com/tensorflow/tensorflow/issues/24412, the TPU support for TensorFlow 2.0 is incomplete yet.

The solution will be to monitor the above issue and wait for the TF 2.0 to be released.