Google ColabでのTPU使い方メモ

各ライブラリで実装する前に, Google Colab のランタイムで TPU を選択する必要があります. 事前に Google Colab メニューバーの Runtime > Change runtime type から Notebook settings を開き,　Hardware accelerator を TPU に変更してください．

以下のサンプルコードでは，簡単な掛け算を TPU を使って実行します．

Tensorflow

import os

import tensorflow as tf

X = tf.placeholder(tf.int32, [5, ])
Y = tf.placeholder(tf.int32, [5, ])

tpu_ops = tf.contrib.tpu.rewrite(tf.multiply, [X, Y])
tpu_address = "grpc://" + os.environ['COLAB_TPU_ADDR']

with tf.Session(tpu_address) as session:
    try:
        session.run(tf.contrib.tpu.initialize_system())
        result = session.run(tpu_ops, {
            X: [1, 2, 3, 4, 5],
            Y: [2, 3, 4, 5, 6]
        })
        # 結果を出力
        print(result)
    finally:
        session.run(tf.contrib.tpu.shutdown_system())

Keras

import os

import numpy as np
import tensorflow as tf

X = tf.keras.Input([5, ])
Y = tf.keras.Input([5, ])

outputs = tf.keras.layers.Multiply()([X, Y])
model = tf.keras.Model([X, Y], outputs)

tpu_address = "grpc://" + os.environ['COLAB_TPU_ADDR']
storategy = tf.contrib.tpu.TPUDistributionStrategy(
    tf.contrib.cluster_resolver.TPUClusterResolver(tpu_address))

tpu_model = tf.contrib.tpu.keras_to_tpu_model(model, storategy)
result = tpu_model.call([
    np.array([1, 2, 3, 4, 5]),
    np.array([2, 3, 4, 5, 6])
])
# 結果を出力
print(result)

tf.contrib.tpu.shutdown_system()