How to use TPUs with Colab

In this section, we show how to use TPUs with Colabs. Just point your browser to https://colab.research.google.com/ and change the runtime from the runtime menu as shown in Figure 9:

Figure 9: Setting TPU as runtime in Colab

Checking whether TPUs are available

First of all, let's check if there is a TPU available by using this simple code fragment that returns the IP address assigned to the TPU. Communication between CPU and TPU happens via grpc:

import os
try:
    device_name = os.environ['COLAB_TPU_ADDR']
    TPU_ADDRESS = 'grpc://' + device_name
    print('Found TPU at: {}'.format(TPU_ADDRESS))
except KeyError:
    print('TPU not found')
Found TPU at: grpc://10.91.166.82:8470

We've confirmed that a TPU is available! Now, we'll continue to explore how we can make use of it.

Loading data with tf.data

Our goal is to implement a simple CNN on MNIST...