Google colablatoryの無料TPU上でtensorflowのKeras APIで実装したCNNを動かしてみる。

久しぶりにDeep Learningを使いたいと思い、兼ねてより気になっていたが今まで使うタイミングがなかったGoogle colabolatoryの無料TPU(※ ただし、12h以内)の上でCNNを動かしてみる。本記事執筆の時点ではTPU対応の深層学習フレームワークはtensorflowのみのようだ。tensorflowは計算グラフがやや書きにくいのだが、近年tensorflowはkeras APIを提供し、抽象化した書き方ができるようになったり、eagar executionというモードではpytorchやchainerのようなDefine by Runの書き方も可能になり、大分扱いやすくなった。そこで本記事ではtensorflowのKeras APIを使ってTPU上でCNNを使ってMNISTの分類タスクの学習をやってみる。

実装

まずは全体のコードを眺めてみましょう。

import tensorflow as tf
from tensorflow.contrib.tpu.python.tpu import keras_support
import os
from keras.datasets import mnist
from keras.utils import np_utils

batch_size = 1024 ＃バッチサイズは大きくする 
num_classes = 10
epochs = 12

img_rows, img_cols = 28, 28

(x_train, y_train), (x_test, y_test) = mnist.load_data() # Kerasでデータセットの取得

x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
input_shape = (img_rows, img_cols, 1)
    
x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train /= 255
x_test /= 255
print('x_train shape:', x_train.shape)
print(x_train.shape[0], 'train samples')
print(x_test.shape[0], 'test samples')

y_train = y_train.astype('int32')
y_test = y_test.astype('int32')
y_train = np_utils.to_categorical(y_train, num_classes)
y_test =  np_utils.to_categorical(y_test, num_classes)

# CNNのモデルを構築
model = tf.keras.Sequential()
model.add(tf.keras.layers.Conv2D(32, kernel_size=(3, 3),
                 activation='relu',
                 input_shape=input_shape))
model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))
model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
model.add(tf.keras.layers.Dropout(0.25))
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(128, activation='relu'))
model.add(tf.keras.layers.Dropout(0.5))
model.add(tf.keras.layers.Dense(num_classes, activation='softmax'))

# TPU用にモデルを変換
tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
model = tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)

model.compile(loss=tf.keras.losses.categorical_crossentropy,
              optimizer=tf.keras.optimizers.RMSprop(),
              metrics=['accuracy'])

model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs,
          verbose=1, validation_data=(x_test, y_test))

# 性能評価
score = model.evaluate(x_test, y_test, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])

TPUを使用するためには

まず、TPUを使用する場合コード内で構築したモデルをTPU用に変換する必要がある。 変換のための処理は以下の4行である。

tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
model = tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)

続いて ランタイム→ ランタイムのタイプを変更を押す

TPUを選択し、保存

これだけでTPUが使用する準備が完了。 あとはコードを実行するだけである。

結果は、１分足らずでMNISTを12epoch回して精度99%を叩き出すことができた。

注意点

• kerasはデータセットの取得と、出力のカテゴリをone-hot encodingする前処理のためのみに使用し、それ以外はtensorflowのKeras APIを用いている。tensorflowもkerasもGoogle Colaboratoryに元々インストールされているので一切のセットアップは要らない。
• TPUを使用する場合バッチサイズは大きくするのがポイントである。実はTPUはバッチサイズが小さいとGPUよりも遅い。batch_size = 128で学習してみたところ 、TPUはGPUに比べて3倍ほど処理が遅いが、batch_size = 1024ではTPUはGPUの２倍ほどの速さとなった。
• 当初optimizerにtensorflowのkeras APIではなく、純粋なkerasのoptimizerメソッドを使用していたら、TPU使用の時のみエラーで動かないケースがあったので注意。