3.3k{icon} {views}
TensorFlow2.0＋Colab TPUでモデルを保存する方法、CPUとTPUで保存した係数を相互運用する方法、TPUを意識したモデルの保存方法を見ていきます。
環境
CPU：Windows 10
TPU: Google Colab TPU
どちらもTensorFlow 2.0.0
CPUでのモデルの保存
訓練ループを書く例を紹介します。
import tensorflow as tf
import tensorflow.keras.layers as layers
import numpy as np
def create_model():
inputs = layers.Input((28, 28))
x = layers.Flatten()(inputs)
x = layers.Dense(128, activation="relu")(x)
x = layers.Dense(10, activation="softmax")(x)
return tf.keras.models.Model(inputs, x)
def load_dataset():
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()
trainset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
trainset = trainset.map(
lambda x, y: (tf.cast(x, tf.float32) /255.0, tf.cast(y, tf.float32))
).shuffle(50000).batch(128)
test_set = tf.data.Dataset.from_tensor_slices((X_test, y_test))
test_set = test_set.map(
lambda x, y: (tf.cast(x, tf.float32) / 255.0, tf.cast(y, tf.float32))
).batch(128)
return trainset, test_set
def train():
trainset, testset = load_dataset()
model = create_model()
optim = tf.keras.optimizers.Adam()
loss_func = tf.keras.losses.SparseCategoricalCrossentropy()
acc = tf.keras.metrics.SparseCategoricalAccuracy()
@tf.function
def train_on_batch(X, y):
with tf.GradientTape() as tape:
y_pred = model(X, training=True)
loss = loss_func(y, y_pred)
grad = tape.gradient(loss, model.trainable_weights)
optim.apply_gradients(zip(grad, model.trainable_weights))
acc.update_state(y, y_pred)
return loss
@tf.function
def validation_on_batch(X, y):
y_pred = model(X, training=False)
loss = loss_func(y, y_pred)
acc.update_state(y, y_pred)
return loss
for epoch in range(10):
acc.reset_states()
trainloss = []
for X, y in trainset:
trainloss.append(train_on_batch(X, y).numpy())
trainloss = np.mean(np.asarray(trainloss))
trainacc = acc.result().numpy()
testloss = []
acc.reset_states()
for X, y in testset:
testloss.append(validation_on_batch(X, y).numpy())
testloss = np.mean(np.asarray(testloss))
print("Epoch : ", epoch + 1, "train_loss : ", trainloss, "train_acc : ", trainacc,
"test_loss : ", testloss, "test_acc : ", acc.result().numpy())
model.save("model.h5")
def inference():
_, testset = load_dataset()
model = tf.keras.models.load_model("model.h5")
acc = tf.keras.metrics.SparseCategoricalAccuracy()
for X, y in testset:
y_pred = model(X, training=False)
acc.update_state(y, y_pred)
print("Test accuracy : ", acc.result().numpy())
trainを実行すると訓練が行われ、
TensorFlowのドキュメントでは、チェックポイントとして保存することが一番最初に出てきますが、TensorFlow2.0.0時点ではTPUとの互換性を考えるとh5形式で保存したほうがよさそうです。なぜなら、TPU環境でチェックポイントとして保存でしようとすると、Google Cloud Storage以外保存できなくなってしまう（エラーが出てローカルファイルに保存できない：詳しくはこちら）という現象があるからです。このエラーはh5形式として保存すると発生しません。
したがって、TPUとの互換性を考えるなら現時点ではh5形式として保存するのが簡単です。
この係数を読みこんで推論する、inferenceを実行すると、
このように、最終エポックと同じ結果が復元されました。
TPUでのモデルの保存
次はTPU環境でのモデルの保存です。Google ColabでTPUを使うための準備をします（以降のTPUでの計算でも同様に必要です）。
コードはこちら
import tensorflow.keras.layers as layers
import numpy as np
def create_model():
inputs = layers.Input((28, 28))
x = layers.Flatten()(inputs)
x = layers.Dense(128, activation="relu")(x)
x = layers.Dense(10, activation="softmax")(x)
return tf.keras.models.Model(inputs, x)
def load_dataset():
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()
trainset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
trainset = trainset.map(
lambda x, y: (tf.cast(x, tf.float32) /255.0, tf.cast(y, tf.float32))
).shuffle(50000).batch(128)
test_set = tf.data.Dataset.from_tensor_slices((X_test, y_test))
test_set = test_set.map(
lambda x, y: (tf.cast(x, tf.float32) / 255.0, tf.cast(y, tf.float32))
).batch(128)
return trainset, test_set
def train():
trainset, testset = load_dataset()
with strategy.scope():
model = create_model()
optim = tf.keras.optimizers.Adam()
loss_func = tf.keras.losses.SparseCategoricalCrossentropy(reduction=tf.keras.losses.Reduction.NONE)
acc = tf.keras.metrics.SparseCategoricalAccuracy()
trainset = strategy.experimental_distribute_dataset(trainset)
testset = strategy.experimental_distribute_dataset(testset)
def train_on_batch(X, y):
with tf.GradientTape() as tape:
y_pred = model(X, training=True)
loss = loss_func(y, y_pred)
loss = tf.reduce_sum(loss, keepdims=True) / 128.0
grad = tape.gradient(loss, model.trainable_weights)
optim.apply_gradients(zip(grad, model.trainable_weights))
acc.update_state(y, y_pred)
return loss
@tf.function
def distributed_train_on_batch(X, y):
loss = strategy.experimental_run_v2(train_on_batch, args=(X, y))
return strategy.reduce(tf.distribute.ReduceOp.SUM, loss, axis=None)
def validation_on_batch(X, y):
y_pred = model(X, training=False)
loss = loss_func(y, y_pred)
loss = tf.reduce_sum(loss, keepdims=True) / 128.0
acc.update_state(y, y_pred)
return loss
@tf.function
def distributed_validation_on_batch(X, y):
loss = strategy.experimental_run_v2(validation_on_batch, args=(X, y))
return strategy.reduce(tf.distribute.ReduceOp.SUM, loss, axis=None)
for epoch in range(10):
acc.reset_states()
trainloss = []
for X, y in trainset:
trainloss.append(distributed_train_on_batch(X, y).numpy())
trainloss = np.mean(np.asarray(trainloss))
trainacc = acc.result().numpy()
testloss = []
acc.reset_states()
for X, y in testset:
testloss.append(distributed_validation_on_batch(X, y).numpy())
testloss = np.mean(np.asarray(testloss))
print("Epoch : ", epoch + 1, "train_loss : ", trainloss, "train_acc : ", trainacc,
"test_loss : ", testloss, "test_acc : ", acc.result().numpy())
model.save("model_tpu.h5")
def inference():
_, testset = load_dataset()
with strategy.scope():
model = tf.keras.models.load_model("model_tpu.h5")
testset = strategy.experimental_distribute_dataset(testset)
acc = tf.keras.metrics.SparseCategoricalAccuracy()
def validation_on_batch(X, y):
y_pred = model(X, training=False)
acc.update_state(y, y_pred)
@tf.function
def distributed_validation_on_batch(X, y):
return strategy.experimental_run_v2(validation_on_batch, args=(X, y))
for X, y in testset:
distributed_validation_on_batch(X, y)
print("Test accuracy : ", acc.result().numpy())
train()を実行すると、（TPUでのMLPはやたら遅いですけど、CNNはとても速いです）
モデルを読み込んで推論します。
TPUでも同様に、精度を再現できました。
TPU→CPUの読み込み
ここからが本番で、TPUで訓練した重みをCPUで読み込んで推論します。実践的にはよくある形ですね。先程の「model_tpu.h5」をダウンロードしてきます。
このようにTPUでの訓練結果が再現されました。
Kerasでcompile→fitで訓練するケースで、h5形式で保存し、かつカスタム損失関数や評価関数がある場合は読み込みで少し引っかかるかもしれません。これはTF1.X系からあるので以下の方法で解決できます。
Kerasでカスタム損失関数を用いたモデルをロードするとValueError
https://blog.shikoan.com/keras-load-custom-loss-model/
CPU→TPUの読み込み
ケースとしては少ないかもしれませんが、さっきの逆をやってみます。ColabにCPUで訓練した「model.h5」をアップロードします。
小数第4位が違いますがほぼ誤差でしょう。
まとめ
h5形式でモデルを保存することで、TPU→CPU、CPU→TPUの相互のやり取りが無事できたということを確認できました。
Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内
