TPU стал бесплатным в Google Colab

Хотя его скорость примерно такая же, как у графического процессора Colab (заблокирована Google?), его все же стоит попробовать, потому что Google может увеличить скорость TPU в будущем.

К сожалению, для использования TPU требуется изменение кода.

Изменение кода

Во-первых, получите имя устройства и адрес TPU.

try:
    device_name = os.environ['COLAB_TPU_ADDR']
    TPU_ADDRESS = 'grpc://' + self.device_name
    print('Found TPU')
except KeyError:
    print('TPU not found')

Во-вторых, измените модель Keras на модель TPU.

tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
       tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)))

Примечание

Чтобы построить модель Keras, следующий объект Keras (включая, но не все) не следует импортировать из keras.*

Они должны быть импортированы из tensorflow.keras.*

например (неправильный)

from keras.models import Sequential, Model
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D

например (правильный)

from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense, Dropout, Flatten
from tensorflow.keras.layers import Conv2D, MaxPooling2D