How can you use TPU from Google Colab in Tensorflow 2.0?

Question

I'm trying to make use of Google Colab to use a Tensor Processing Unit (TPU) to train a neural network. Tensorflow has just come out with a major release, 2.0, so I am trying to do this within Tensorflow 2.0. I have tried following three guides, but all were written for Tensorflow 1.14- and fail with Tensorflow 2.0:

1) Following the guide TPUs in Colab , I get the error:

AttributeError: module 'tensorflow' has no attribute 'Session'

(from the reference: with tf.Session(tpu_address) as session: )

2) Following the guide Simple Classification Model using Keras on Colab TPU , I get the same error

3) Following the guide cloud_tpu_custom_training , I get the error:

AttributeError: module 'tensorflow' has no attribute 'contrib'

(from the reference: resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu=TPU_WORKER) )

Does anyone have an example of using a TPU to train a neural network in Tensorflow 2.0?

Edit: This issue also appears to have been raised on github: InvalidArgumentError: Unable to find a context_id matching the specified one #1

Answer 1

Firstly the code given in the tutorial is not 2.x compatible

1. You need to choose runtime as TPU in colab to execute code in TPU

2. For the error

   AttributeError: module 'tensorflow' has no attribute 'Session'

   you need to use tf.compat.v1.Session() as tf.session is deprecated.

3. In place of tf.contrib.cluster_resolver please use tf.distribute.cluster_resolver

Please refer Tensorflow Addon-repo to convert code from 1.x to 2.x compatible.

Answer 2

Finally support has been added for TPUs in Tensorflow 2.1.0 (as of Jan 8, 2020). From the release notes here https://github.com/tensorflow/tensorflow/releases/tag/v2.1.0 :

Experimental support for Keras.compile, .fit, .evaluate, and.predict is available for Cloud TPUs, Cloud TPU, for all types of Keras models (sequential, functional and subclassing models).

The tutorial is available here: https://www.tensorflow.org/guide/tpu

For completeness, I'll add the walkthrough here:

1. Go to Google Colab and create a new Python 3 Notebook here: https://colab.research.google.com/
2. In the toolbar, click Runtime / Change runtime type, then choose "TPU" under Hardware accelerator.
3. Copy and paste the below code into the notebook and click run cell (the play button).

from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf
import os
import tensorflow_datasets as tfds

# Distribution strategies
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

# MNIST model
def create_model():
  return tf.keras.Sequential(
      [tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
       tf.keras.layers.Flatten(),
       tf.keras.layers.Dense(128, activation='relu'),
       tf.keras.layers.Dense(10)])

# Input datasets
def get_dataset(batch_size=200):
  datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True,
                             try_gcs=True)
  mnist_train, mnist_test = datasets['train'], datasets['test']

  def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255.0

    return image, label

  train_dataset = mnist_train.map(scale).shuffle(10000).batch(batch_size)
  test_dataset = mnist_test.map(scale).batch(batch_size)

  return train_dataset, test_dataset

# Create and train a model
strategy = tf.distribute.experimental.TPUStrategy(resolver)
with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

train_dataset, test_dataset = get_dataset()

model.fit(train_dataset,
          epochs=5,
          validation_data=test_dataset,steps_per_epoch=50)

Note that when I run the code from the tensorflow tutorial as-is I get the below error. I've corrected this by adding the steps_per_epoch parameter in model.fit()

ValueError: Number of steps could not be inferred from the data, please pass the steps_per_epoch argument.

Answer 3

Before running the code,

Go to,

Edit --> Notebook Settings

Under that select

Hardware Accelerator --> TPU

Answer 4

Although TF2.0 has been released, but the colab version hasn't been updated. We are working on the colab update and will make sure this is ready with the 2.1 release

See link for more details: LINK

Answer 5

The upgraded version of "tf" will resolve the said issue.

.pip install tensorflow==2.7.0

Answer 6

Tensorflow 2.0 isn't really backwards compatible with Tensorflow 1.X code. There are quite a few changes to how Tensorflow works between these versions, so I highly recommend reading the official guide on how to migrate your code:

https://www.tensorflow.org/guide/migrate#estimators

I will say, the automatic conversion script, though technically successful, only changed my code to the compatibility versions of the Tensorflow 1.X code - if you want to use any actual Tensorflow 2.0 features, you'll probably need to change the code manually.