GoogleColabでTPUを使う（TensorFlow 2.8）

Google ColaboratoryのランタイムタイプのハードウェアアクセラレータでGPU以外にも「TPU」が選べる。 TPUを使用する場合にコードの変更が必要だった。

あらまし

書籍「AlphaZero 人工知能プログラミング実践入門」の３章３節「畳み込みニューラルネットワークで画像分類」ではCIFAR-10の画像分類を行うが、 その後半でGoogleColabのランタイムタイプをTPUにする際の変更点が示されている。 指示通りに動かしてみようとしたところバージョンの違いか、エラーが出てしまう。 これを動くようにしてみた。

本が想定しているバージョン（執筆時点：2019年5月）：

• Python 3.6.7
• tensorflow 1.13.1

動かしたバージョン：

• tensorflow 2.8.0

修正点

本には作成したモデルを tf.contrib.tpu.keras_to_tpu_model でTPU用に変換するよう書かれているが、contrib は統合され、コンバートする関数はなくなったぽい。 TPUs in Colab - Colaboratory を参考にする。

TPUの初期化

%tensorflow_version 2.x
import tensorflow as tf
print("Tensorflow version " + tf.__version__)

try:
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection
  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])
except ValueError:
  raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')

tf.config.experimental_connect_to_cluster(tpu)
tf.tpu.experimental.initialize_tpu_system(tpu)
tpu_strategy = tf.distribute.TPUStrategy(tpu)

• 「TPUs in Colab」のままだと WARNING:absl:`tf.distribute.experimental.TPUStrategy` is deprecated, please use the non experimental symbol `tf.distribute.TPUStrategy` instead. というワーニングが出ているので、 tf.distribute.TPUStrategy に変更する

モデルの作成

モデルの作成は、直接行う代わりに with にストラテジのスコープを指定した状態で行う：

with tpu_strategy.scope():
  # モデルの作成
  model = Sequential()
  ...

あとがき

以上の修正でTPUで動くようになった。

• 全ソース
• GPU 207.707秒に対して、TPU 178.723秒でした（GPUに対し1.16倍高速）
  • より複雑なモデルなら効果は大きいかもしれない
  • 本にはGPU 6,800秒、TPU 225秒と書いてあり、当時とはGPUのスペックが変わったのかも
• TPUの場合はバッチサイズは tpu_strategy.num_replicas_in_sync から決定するほうがいい？ 下手に変更するとAccuracyに影響が出る。
• TPUの場合には推論のテスト時に『「学習データ数」 と「バッチサイズ」は「TPUコアの数」で割り切れる必要があります』として10から16に変更するよう書いてあるが、そうしなくても動いた
• TensorFlowのバージョンが変わるたびに修正が必要なのは困ったもんだ…
  • config.experimental_connect_to_cluster とか、また tpu には experimental しかないのでこれも後々修正になるんだろう
  • オプティマイザの Adam に渡す引数が lr だとワーニングで learning_rate に修正する必要がある
• TPUを使用する  |  TensorFlow Core