必要なことまとめ
- ランタイムで「TPU」を選択する
- kerasではなくtensorflow.kerasを使う
- modelをTPU用のモデルに変換する
- TPUモデルではpredictができないので確認はCPUモデルに戻して行う
Google ColabでTPU使うのは、こちらの記事が詳しいです。自分もこちらを参考に試しました。
Google ColaboratoryのTPUを試してみる
またCholletさんが紹介されてた公式？のノートがあり、必要なコードがまとまってるのでおすすめです。
Keras Fashion MNIST and TPUs - Colaboratory
kerasではなくtensorflow.kerasを使う
下記のように「keras」からのものを使うとTPU未対応のためエラーが出て動きません。
from keras.models import Sequential model = Sequential()
「keras」を「tensorflow.keras」にすべて置き換えてやる必要があります。
from tensorflow.keras.models import Sequential model = Sequential()
modelをTPU用のモデルに変換する
modelは普通にそのまま作れますが、その作ったモデルをkeras_to_tpu_modelを使ってTPU用のモデルに変換して、compileやfitします。
また、model生成にはOSからTPU情報取得して渡す必要があります。
import tensorflow as tf import os 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_model.compile(...) tpu_model.fit(...)
ランタイムでGPU選択しただけでGPUが使えるため、ついTPUを選択しただけで早くなるのかと思ってしまいがちですが、変換しないとCPUのみになるのでめちゃ遅いので注意です。
TPUモデルではpredictができないので確認はCPUモデルに戻して行う
tpu_model.predict()がないため、学習後の重みを使って出力させることができません。
どうするのかな？と思ったら、CPUモデルに戻してやってそれでpredictする必要があるようです。
cpu_model = tpu_model.sync_to_cpu() cpu_model.predict(...)
ColabのGPUとTPUランタイムでの速度差
Keras Fashion MNIST and TPUs - Colaboratory で使われている Conv2D x3 + Dense x2 のサンプルでは
GPU: 235秒
TPU: 89秒
となり、2.6倍程度の速度でした。
が、TPUはバッチサイズをもっと大きくすることができ、またネットワークサイズをGPUよりもだいぶ大きくして動かすことができます。
ネットワークサイズが大きくなるほど差が大きくなる感じです。
