【tensorflow2.x 入門】@tf.function で学習速度が上がるか試してみる

tf2,x 系に最近入門してみて、subclassing API まわりの書き方を勉強しています。

学習ループの中で、逆伝播を行う tf.GradientTape() 以下の箇所が計算負荷が大きいが、@tf.function を追加することで、関数を計算グラフに変換してより早く実行してくれる。ということだそうなので、実際に追加してみて速度がどの程度変わるのかを試してみました。

小さなデータセットでは速度の差をあまり実感できなさそうなので、今回は food101 データセットを利用します。

データセットについてはこちらに書いてあります。512×512サイズの画像で、101クラス分類です。

まずは、tensorflow_datasets から food101 をダウンロードして、224×224 にリサイズします。メモリの問題で、バッチサイズは32にして、データセットを作成します。

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)

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 の学習済みモデルの特徴抽出器の部分のみを取り出します。重みは固定せず学習させます。

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

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 なしで学習とテストの関数を作成します。

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)

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 エポックで何秒かかったか測ります。

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
'''

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 エポック回します。

@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
'''