【Python】TensorFlowでTPUを使う方法

TPUを使えばTensorFlowの学習を高速化できます。

使用するコードは以下の記事で解説しているものです。

サンプルコード

import tensorflow as tf
try:
    # TPU detection. No parameters necessary if TPU_NAME environment variable is
    # set: this is always the case on Kaggle.
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
    print('Running on TPU ', tpu.master())
except ValueError:
    tpu = None

if tpu:
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    strategy = tf.distribute.TPUStrategy(tpu)
else:
    # Default distribution strategy in Tensorflow. Works on CPU and single GPU.
    strategy = tf.distribute.get_strategy()

AUTO = tf.data.experimental.AUTOTUNE
REPLICAS = strategy.num_replicas_in_sync
print("REPLICAS: ", REPLICAS)

import numpy as np
#import tensorflow as tf

seed = 0
np.random.seed(seed)
tf.random.set_seed(seed)

data = tf.keras.datasets.cifar10.load_data()

(X_train, y_train), (X_valid, y_valid) = data

y_train = tf.one_hot(y_train.ravel(), depth = 10)
y_valid = tf.one_hot(y_valid.ravel(), depth = 10)

print(X_train.shape, y_train.shape, X_valid.shape, y_valid.shape)

import matplotlib.pyplot as plt
plt.imshow(X_train[0])
plt.title(y_train[0].numpy().astype("int8").tolist())
plt.show()

batch_size = 32 * REPLICAS
print(batch_size)

train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
train_dataset = train_dataset.batch(batch_size)
train_dataset = train_dataset.shuffle(1024 * 50)

valid_dataset = tf.data.Dataset.from_tensor_slices((X_valid, y_valid))
valid_dataset = valid_dataset.batch(batch_size * 2)

def augmentation(image, label):
    image = tf.image.random_flip_left_right(image)
    return image, label

train_dataset = train_dataset.map(augmentation)

import tensorflow.keras.layers as L
import tensorflow.keras.models as M

def get_model(image_size, num_classes):
    inputs = L.Input(shape = image_size)
    x = tf.keras.applications.efficientnet.EfficientNetB0(include_top = False)(inputs)
    x = L.Flatten()(x)
    outputs = L.Dense(units = num_classes, activation = "softmax")(x)

    model = M.Model(inputs = inputs, outputs = outputs)
    model.compile(
        optimizer = tf.keras.optimizers.Adam(learning_rate = 0.001),
        loss = "categorical_crossentropy"
    )
    return model

tf.keras.backend.clear_session()
with strategy.scope():
    model = get_model(X_train.shape[-3:], y_train.shape[-1])
model.summary()

checkpoint = tf.keras.callbacks.ModelCheckpoint(
    "best_weight.h5",
    monitor = "val_loss",
    direction = "min",
    save_best_only = True,
    save_weights_only = True
)

epochs = 10

scheduler = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate = 0.001, decay_steps = epochs,
)

history = model.fit(
    train_dataset,
    validation_data = valid_dataset,
    epochs = epochs,
    callbacks = [
        checkpoint,
        tf.keras.callbacks.LearningRateScheduler(scheduler)
    ]
)

import matplotlib.pyplot as plt
plt.plot(history.history["loss"], label = "train")
plt.plot(history.history["val_loss"], label = "valid")
plt.legend()
plt.show()

from sklearn.metrics import accuracy_score, precision_score, recall_score
model.load_weights("best_weight.h5")
pred = model.predict(X_valid, verbose = 1).argmax(axis = 1)
truth = y_valid.numpy().argmax(axis = 1)
accuracy = accuracy_score(truth, pred)
precision = precision_score(truth, pred, average = "macro")
recall = recall_score(truth, pred, average = "macro")
print("ACC:", accuracy.round(3), "PRE:", precision.round(3), "REC:", recall.round(3))

データセットCIFAR10の画像を分類するモデルです。

TPUで計算する方法

import tensorflow as tf
try:
    # TPU detection. No parameters necessary if TPU_NAME environment variable is
    # set: this is always the case on Kaggle.
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
    print('Running on TPU ', tpu.master())
except ValueError:
    tpu = None

if tpu:
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    strategy = tf.distribute.TPUStrategy(tpu)
else:
    # Default distribution strategy in Tensorflow. Works on CPU and single GPU.
    strategy = tf.distribute.get_strategy()

AUTO = tf.data.experimental.AUTOTUNE
REPLICAS = strategy.num_replicas_in_sync
print("REPLICAS: ", REPLICAS)

重要な設定は上記コードの部分です。

Kaggleでよく使われる定型文をそのまま持ってきました。

その次にbatch_sizeの設定です。

batch_size = 32 * REPLICAS

REPLICASの数だけ乗算してしていします。

例えばCPU,GPUで計算するときはREPLICAS = 1になっているはずです。

なので上記コードではbatch_size = 32 x 1となります。

ですがTPUを使うとREPLICAS = 8となるので、batch_size = 32 x 8です。

これだけバッチサイズを大きくとれるのがメリットの１つです。

最後にモデルを呼び出す部分です。

tf.keras.backend.clear_session()
with strategy.scope():
    model = get_model(X_train.shape[-3:], y_train.shape[-1])
model.summary()

strategy.scopeの直下でモデルを呼び出します。

何しているのか正直わかっていませんが、こうすることでTPUを使用したモデルになります。

計算速度

CPU

CPUで冒頭のサンプルコードを実行すると１epochが10分くらいかかります。

待ってられないので結果の画像は省略します。

GPU

GPUを使うとCPUでの計算から圧倒的に速くなります。

下図はKaggleのGPU(P100)で計算した例です。

だいたい１epochが１分半くらいですね。

TPU

最後にTPUを使った結果です。

GoogleColabのTPUを使いました。

なので、実はKaggleのTPUを使うほうが圧倒的に速いです。

ですがKaggleのTPUが人気過ぎてよく順番待ちになるので、GoogleColabの方を使いました。

１epochが9秒くらいで圧倒的に速いです。

もちろんバッチサイズが８倍とれていることが要因の１つです。

まとめ

今回はTensorFlowでTPUを使う方法を解説しました。

こんな感じです。