【Keras】GoogleColabで簡単にTPUを試す方法

【重要】tensorflow から keras をインポートします。

from tensorflow import keras as keras
import tensorflow as tf
import numpy as np
import os

ダミーデータを作成

X_train = np.random.rand(1000,10)
y_train = np.random.randint(3, size=1000)
y_train = keras.utils.to_categorical(y_train)

print(X_train.shape)
print(y_train.shape)

kerasで適当なmodelを作成

model = keras.models.Sequential()
model.add(keras.layers.Dense(64,input_shape=(10,),activation='relu'))
model.add(keras.layers.Dense(3,activation='softmax'))
model.compile(loss='categorical_crossentropy',optimizer =keras.optimizers.RMSprop(),metrics=['categorical_accuracy'])
model.summary()

※ ここから重要
kerasのmodelをTPU用のモデルに変換

tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(
            tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
    )
)

TPU用のモデルで学習

tpu_model.fit(X_train,y_train,epochs=10,batch_size=10)

上記のプログラムはGoogleColabeでこちら から実行できます。