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です。
これだけバッチサイズを大きくとれるのがメリットの1つです。
最後にモデルを呼び出す部分です。
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で冒頭のサンプルコードを実行すると1epochが10分くらいかかります。
待ってられないので結果の画像は省略します。
GPU
GPUを使うとCPUでの計算から圧倒的に速くなります。
下図はKaggleのGPU(P100)で計算した例です。


だいたい1epochが1分半くらいですね。
TPU
最後にTPUを使った結果です。
GoogleColabのTPUを使いました。
KaggleのTPU:version3
なので、実はKaggleのTPUを使うほうが圧倒的に速いです。
ですがKaggleのTPUが人気過ぎてよく順番待ちになるので、GoogleColabの方を使いました。


1epochが9秒くらいで圧倒的に速いです。
もちろんバッチサイズが8倍とれていることが要因の1つです。
まとめ
今回はTensorFlowでTPUを使う方法を解説しました。
・バッチサイズを大きくする
・strategy.scope直下でモデルを呼び出す
こんな感じです。
コメント