TPU is now free in Google Colab
Although its speed is just similar to Colab GPU (locked by Google?), it is still worth to try it, because Google may increase TPU speed in the future.
To use TPU, unfortunately, code change is required.
Code change
Firstly, get device name and TPU address
try:
device_name = os.environ['COLAB_TPU_ADDR']
TPU_ADDRESS = 'grpc://' + self.device_name
print('Found TPU')
except KeyError:
print('TPU not found')
Secondly, change Keras model into TPU model
tf.contrib.tpu.keras_to_tpu_model(
model,
strategy=tf.contrib.tpu.TPUDistributionStrategy(
tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)))
Note
To build Keras model, the following Keras object (include but not all) should not be imported from keras.*
They should be imported from tensorflow.keras.*
e.g. (wrong)
from keras.models import Sequential, Model
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D
e.g. (correct)
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense, Dropout, Flatten
from tensorflow.keras.layers import Conv2D, MaxPooling2D