tf2,x 系に最近入門してみて、subclassing API まわりの書き方を勉強しています。
学習ループの中で、逆伝播を行う tf.GradientTape() 以下の箇所が計算負荷が大きいが、@tf.function を追加することで、関数を計算グラフに変換してより早く実行してくれる。ということだそうなので、実際に追加してみて速度がどの程度変わるのかを試してみました。
小さなデータセットでは速度の差をあまり実感できなさそうなので、今回は food101 データセットを利用します。
データセットについてはこちらに書いてあります。512×512サイズの画像で、101クラス分類です。
まずは、tensorflow_datasets から food101 をダウンロードして、224×224 にリサイズします。メモリの問題で、バッチサイズは32にして、データセットを作成します。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
import numpy as np import tensorflow as tf from tensorflow.keras import datasets, layers, models, optimizers, losses, metrics import tensorflow_datasets as tfds datasets = tfds.load("food101", as_supervised=True) train, test = datasets['train'], datasets['validation'] def scale(image, label): image = tf.cast(image, tf.float32) image /= 255.0 image = tf.image.resize(image, [224, 224]) return image, label batch_size = 32 train_ds = train.map(scale).shuffle(10000, seed=0).batch(batch_size) test_ds = test.map(scale).batch(batch_size) |
tensorflow-hub から、mobilenet_v2 の学習済みモデルの特徴抽出器の部分のみを取り出します。重みは固定せず学習させます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
import tensorflow_hub as hub # モデル class Net(models.Model): def __init__(self, output_dim=101): super().__init__() self.model = models.Sequential([ hub.KerasLayer( "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/2", input_shape=(224, 224, 3), trainable=True ), layers.Dense(output_dim, activation="softmax") ]) def call(self, x): return self.model(x) net = Net() criterion = losses.SparseCategoricalCrossentropy() optimizer = optimizers.SGD(learning_rate=1e-3, momentum=0.9, nesterov=True) train_loss = metrics.Mean() train_acc = metrics.SparseCategoricalAccuracy() test_loss = metrics.Mean() test_acc = metrics.SparseCategoricalAccuracy() epochs = 1 |
まずは、@tf.function なしで学習とテストの関数を作成します。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
def train(x, t): with tf.GradientTape() as tape: preds = net(x) loss = criterion(t, preds) grads = tape.gradient(loss, net.trainable_variables) optimizer.apply_gradients(zip(grads, net.trainable_variables)) train_loss(loss) train_acc(t, preds) def test(x, t): preds = net(x) loss = criterion(t, preds) test_loss(loss) test_acc(t, preds) |
スタートタイムとエンドタイムの差で、1 エポックで何秒かかったか測ります。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
import time start = time.time() print('start {}'.format(start)) for epoch in range(epochs): for batch in train_ds: x, t = batch train(x, t) for batch in test_ds: x, t = batch test(x, t) train_loss.reset_states() train_acc.reset_states() test_loss.reset_states() test_acc.reset_states() end = time.time() print('end {}'.format(end)) print("@tf.function なし：{}".format(end - start)) ''' start 1599118735.4095879 end 1599119590.4397452 @tf.function なし：855.030157327652 ''' |
約 855 秒ですので、14 分ちょっとですね。
では train 関数に @tf.function をつけて再度 1 エポック回します。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
@tf.function def train(x, t): with tf.GradientTape() as tape: preds = net(x) loss = criterion(t, preds) grads = tape.gradient(loss, net.trainable_variables) optimizer.apply_gradients(zip(grads, net.trainable_variables)) train_loss(loss) train_acc(t, preds) return loss def test(x, t): preds = net(x) loss = criterion(t, preds) test_loss(loss) test_acc(t, preds) start = time.time() print('start {}'.format(start)) for epoch in range(epochs): for batch in train_ds: x, t = batch train(x, t) for batch in test_ds: x, t = batch test(x, t) train_loss.reset_states() train_acc.reset_states() test_loss.reset_states() test_acc.reset_states() end = time.time() print('end {}'.format(end)) print("@tf.function あり：{}".format(end - start)) ''' start 1599119591.4841387 end 1599120187.5748258 @tf.function あり：596.0906870365143 ''' |
596 秒！およそ 10 分なので、1/3 ほど学習時間が短縮されました。
▼公式ドキュメントの解説
https://www.tensorflow.org/tutorials/customization/performance?hl=ja
おわり
