Help us understand the problem. What is going on with this article?

【TF2.0応用編】tf.data.DatasetでDataAugmentationをクッソ早くする

はじめに

この記事はTF2.0AdventCalendarで、tf.data.Datasetを使うことによる副作用的な物を検証したものになります。
実際、tf.data.Datasetを使ってmodel.fit_generatorがどれくらいの速度になるのかを検証しました。

CIFAR10で検証

データセットを用意しましょう。

データセットの用意
import tensorflow as tf
import tensorflow.keras as keras
import matplotlib.pyplot as plt
import sklearn
import numpy as np
from tqdm import tqdm

(tr_x,tr_y),(te_x,te_y)=keras.datasets.cifar10.load_data()
tr_x, te_x = tr_x/255.0, te_x/255.0
tr_y, te_y = tr_y.reshape(-1,1), te_y.reshape(-1,1)

tr_ds = tf.data.Dataset.from_tensor_slices((tr_x,tr_y)).batch(128)

for img,label in tr_ds:
    print(img.shape,label.shape)
    break
結果
(128, 32, 32, 3) (128, 1)

正常に出力されているようです。

モデルに入れる

検証用のモデルを作りましょう。

テキトーに重いモデルを作る
model = keras.models.Sequential()
model.add(keras.layers.Convolution2D(32,3,padding="same",activation="relu",input_shape=(32,32,3)))
model.add(keras.layers.Convolution2D(32,3,padding="same",activation="relu"))
model.add(keras.layers.Convolution2D(32,3,padding="same",activation="relu"))
model.add(keras.layers.MaxPooling2D())
model.add(keras.layers.Convolution2D(128,3,padding="same",activation="relu"))
model.add(keras.layers.Convolution2D(128,3,padding="same",activation="relu"))
model.add(keras.layers.Convolution2D(128,3,padding="same",activation="relu"))
model.add(keras.layers.Convolution2D(128,3,padding="same",activation="relu"))
model.add(keras.layers.Convolution2D(128,3,padding="same",activation="relu"))
model.add(keras.layers.MaxPooling2D())
model.add(keras.layers.Convolution2D(256,3,padding="same",activation="relu"))
model.add(keras.layers.Convolution2D(256,3,padding="same",activation="relu"))
model.add(keras.layers.Convolution2D(256,3,padding="same",activation="relu"))
model.add(keras.layers.GlobalAveragePooling2D())
model.add(keras.layers.Dense(1000,activation="relu"))
model.add(keras.layers.Dense(128,activation="relu"))
model.add(keras.layers.Dense(10,activation="softmax"))
model.compile(loss="sparse_categorical_crossentropy",metrics=["accuracy"])
model.summary()
結果
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d (Conv2D)              (None, 32, 32, 32)        896       
_________________________________________________________________
conv2d_1 (Conv2D)            (None, 32, 32, 32)        9248      
_________________________________________________________________
conv2d_2 (Conv2D)            (None, 32, 32, 32)        9248      
_________________________________________________________________
max_pooling2d (MaxPooling2D) (None, 16, 16, 32)        0         
_________________________________________________________________
conv2d_3 (Conv2D)            (None, 16, 16, 128)       36992     
_________________________________________________________________
conv2d_4 (Conv2D)            (None, 16, 16, 128)       147584    
_________________________________________________________________
conv2d_5 (Conv2D)            (None, 16, 16, 128)       147584    
_________________________________________________________________
conv2d_6 (Conv2D)            (None, 16, 16, 128)       147584    
_________________________________________________________________
conv2d_7 (Conv2D)            (None, 16, 16, 128)       147584    
_________________________________________________________________
max_pooling2d_1 (MaxPooling2 (None, 8, 8, 128)         0         
_________________________________________________________________
conv2d_8 (Conv2D)            (None, 8, 8, 256)         295168    
_________________________________________________________________
conv2d_9 (Conv2D)            (None, 8, 8, 256)         590080    
_________________________________________________________________
conv2d_10 (Conv2D)           (None, 8, 8, 256)         590080    
_________________________________________________________________
global_average_pooling2d (Gl (None, 256)               0         
_________________________________________________________________
dense (Dense)                (None, 1000)              257000    
_________________________________________________________________
dense_1 (Dense)              (None, 128)               128128    
_________________________________________________________________
dense_2 (Dense)              (None, 10)                1290      
=================================================================
Total params: 2,508,466
Trainable params: 2,508,466
Non-trainable params: 0
_________________________________________________________________

いい感じに重そうですね。

データセットを入れてみる。

まずは通常の入れ方

通常の入れ方
%%time
model.fit(tr_x,tr_y,batch_size=128,epochs=5)
time
Train on 50000 samples
Epoch 1/5
50000/50000 [==============================] - 7s 145us/sample - loss: 0.6238 - accuracy: 0.7833
Epoch 2/5
50000/50000 [==============================] - 7s 135us/sample - loss: 0.5584 - accuracy: 0.8104
Epoch 3/5
50000/50000 [==============================] - 7s 136us/sample - loss: 0.4936 - accuracy: 0.8303
Epoch 4/5
50000/50000 [==============================] - 7s 139us/sample - loss: 0.4467 - accuracy: 0.8479
Epoch 5/5
50000/50000 [==============================] - 7s 136us/sample - loss: 0.4016 - accuracy: 0.8631
CPU times: user 35.5 s, sys: 4.15 s, total: 39.6 s
Wall time: 35.1 s

nvtopでの画像がこちら
スクリーンショット 2019-12-05 22.01.58.png

結構いい感じにGPUを使ってるように見えます。

tf.data.Datasetで入れる

次にtf.data.Datasetでデータを入れましょう。

datasetを入れる
%%time
model.fit_generator(tr_ds,epochs=5)
結果
Epoch 1/5
391/391 [==============================] - 12s 31ms/step - loss: 2.3748 - accuracy: 0.1219
Epoch 2/5
391/391 [==============================] - 12s 31ms/step - loss: 1.9375 - accuracy: 0.2519
Epoch 3/5
391/391 [==============================] - 12s 31ms/step - loss: 1.6902 - accuracy: 0.3601
Epoch 4/5
391/391 [==============================] - 12s 31ms/step - loss: 1.5300 - accuracy: 0.4300
Epoch 5/5
391/391 [==============================] - 12s 30ms/step - loss: 1.3649 - accuracy: 0.4988
CPU times: user 58.7 s, sys: 1.63 s, total: 1min
Wall time: 1min 1s

あら。倍かかっちゃってますね。
nvtopを見てみましょう

スクリーンショット 2019-12-05 22.08.59.png

あまり実力を発揮できていませんね💦

keras.preprocessing.image.ImageDataGenerator()でチャレンジ

めちゃめちゃ遅かったこいつ、keras.preprocessing.image.ImageDataGenerator()はどうでしょうか

%%time
datagen = keras.preprocessing.image.ImageDataGenerator()
model.fit_generator(datagen.flow(tr_x, tr_y, batch_size=128),
                    steps_per_epoch=len(tr_x) / 128, epochs=5)
結果
Epoch 1/5
391/390 [==============================] - 16s 40ms/step - loss: 14.5061 - accuracy: 0.1000
Epoch 2/5
391/390 [==============================] - 16s 40ms/step - loss: 14.5057 - accuracy: 0.1000
Epoch 3/5
391/390 [==============================] - 16s 40ms/step - loss: 14.5065 - accuracy: 0.1000
Epoch 4/5
391/390 [==============================] - 16s 40ms/step - loss: 14.5067 - accuracy: 0.1000
Epoch 5/5
391/390 [==============================] - 16s 40ms/step - loss: 14.5067 - accuracy: 0.1000
CPU times: user 1min 3s, sys: 1.37 s, total: 1min 4s
Wall time: 1min 18s

さらに遅い・・・通常モードですらこんな重いんですね。
スクリーンショット 2019-12-05 22.12.02.png

これはtf.data.Datasetが重いってこと?

実はtf.data.Datasetが重いというより、Kerasのfit_gerenatorが重いっぽいんですよね。
tf.data.Datasetはフルバッチにすることもできるため、それでfitさせてみましょう。

フルバッチDataset
%%time
tr_ds = tf.data.Dataset.from_tensor_slices((tr_x,tr_y)).batch(tr_x.shape[0]).repeat(5)
for x,y in tr_ds:
    model.fit(x,y,batch_size=128)
結果
Train on 50000 samples
50000/50000 [==============================] - 9s 178us/sample - loss: 2.4252 - accuracy: 0.1054
Train on 50000 samples
50000/50000 [==============================] - 7s 134us/sample - loss: 2.1042 - accuracy: 0.2025
Train on 50000 samples
50000/50000 [==============================] - 7s 137us/sample - loss: 1.7758 - accuracy: 0.3369
Train on 50000 samples
50000/50000 [==============================] - 7s 134us/sample - loss: 1.5804 - accuracy: 0.4166
Train on 50000 samples
50000/50000 [==============================] - 7s 136us/sample - loss: 1.3802 - accuracy: 0.4969
CPU times: user 38.2 s, sys: 6.05 s, total: 44.2 s
Wall time: 39.6 s

あら、めっちゃ早いじゃないっすか!nvtopもこの通り
スクリーンショット 2019-12-05 22.18.01.png

それ、tf.data.Dataset使う意味ある?

意味あります。なぜって・・・それはDataAugmentationに決まってるじゃないっすか!!
通常、フルバッチデータに個々にランダムなTransformをして、それをepoch分連続して出力なんて骨が折れる作業です。
でも、先日のAdventCalendarを思い出してください。こいつはそれがフルバッチでできるんです。素敵すぎませんか?

なら、速度検証だ!

やることは以下の2つです。

  • random rotate
  • random flip(h,v)

速度検証していきましょう。

ImageDataGeneratorでやる

データ確認

ImageDataGen
%%time
datagen = keras.preprocessing.image.ImageDataGenerator(
    horizontal_flip=True,
    vertical_flip=True,
    rotation_range=30,
)

labels = np.array([
    'airplane',
    'automobile',
    'bird',
    'cat',
    'deer',
    'dog',
    'frog',
    'horse',
    'ship',
    'truck'])
plt.figure(figsize=(10,10),facecolor="white")
for b_img,b_label in datagen.flow(tr_x,tr_y,batch_size=25):
    for i, img,label in zip(range(25),b_img,b_label):
        plt.subplot(5,5,i+1)
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        plt.imshow(img)
        plt.xlabel(labels[label])
    break
plt.show()

CIFAR10-ImageDataGen.png

速度検証

speedtest
%%time
datagen = keras.preprocessing.image.ImageDataGenerator(
    horizontal_flip=True,
    vertical_flip=True,
    rotation_range=30,
)
model.fit_generator(datagen.flow(tr_x, tr_y, batch_size=128),
                    steps_per_epoch=len(tr_x) / 128, epochs=5)
結果
Epoch 1/5
391/390 [==============================] - 18s 47ms/step - loss: 1.5626 - accuracy: 0.4296
Epoch 2/5
391/390 [==============================] - 18s 46ms/step - loss: 1.4149 - accuracy: 0.4894
Epoch 3/5
391/390 [==============================] - 18s 46ms/step - loss: 1.3278 - accuracy: 0.5236
Epoch 4/5
391/390 [==============================] - 18s 47ms/step - loss: 1.2493 - accuracy: 0.5516
Epoch 5/5
391/390 [==============================] - 18s 46ms/step - loss: 1.1884 - accuracy: 0.5787
CPU times: user 2min, sys: 4.74 s, total: 2min 5s
Wall time: 1min 30s

うーむ、遅いですね。

フルバッチtf.data.Datasetでやる

まずデータ確認をしましょう

データ確認
import tensorflow_addons as tfa
@tf.function
def rotate_tf(image,label):
    if image.shape.__len__() ==4:

        random_angles = tf.random.uniform(shape = (tf.shape(image)[0], ), minval = -30*np
        .pi / 180, maxval = 30*np.pi / 180)
    if image.shape.__len__() ==3:
        random_angles = tf.random.uniform(shape = (), minval = -30*np
        .pi / 180, maxval = 30*np.pi / 180)

    return tfa.image.rotate(image,random_angles),label

@tf.function
def flip_left_right(image,label):
    return tf.image.random_flip_left_right(image),label

@tf.function
def flip_up_down(image,label):
    return tf.image.random_flip_up_down(image),label

tr_ds = tf.data.Dataset.from_tensor_slices((tr_x,tr_y)).shuffle(40000)
tr_ds = tr_ds.batch(25).map(flip_up_down).map(flip_left_right).map(rotate_tf)

plt.figure(figsize=(10,10),facecolor="white")
for b_img,b_label in tr_ds:
    for i, img,label in zip(range(25),b_img,b_label):
        plt.subplot(5,5,i+1)
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        plt.imshow(img)
        plt.xlabel(labels[label])
    break
plt.show()

CIFAR10-tr_ds_gen.png

うーん、なんか荒い?

これでフルバッチでやってみましょう。

速度検証

speedtest
tr_ds = tf.data.Dataset.from_tensor_slices((tr_x,tr_y)).shuffle(40000)
tr_ds = tr_ds.batch(tr_x.shape[0]).map(flip_up_down).map(flip_left_right).map(rotate_tf).repeat(5)

for img,label in tr_ds:
    model.fit(x=img,y=label,batch_size=128)
結果
Train on 50000 samples
50000/50000 [==============================] - 8s 153us/sample - loss: 2.3030 - accuracy: 0.1447
Train on 50000 samples
50000/50000 [==============================] - 7s 133us/sample - loss: 2.0370 - accuracy: 0.2188
Train on 50000 samples
50000/50000 [==============================] - 7s 133us/sample - loss: 1.8615 - accuracy: 0.2985
Train on 50000 samples
50000/50000 [==============================] - 7s 133us/sample - loss: 1.7357 - accuracy: 0.3570
Train on 50000 samples
50000/50000 [==============================] - 7s 133us/sample - loss: 1.6400 - accuracy: 0.3976
CPU times: user 1min 42s, sys: 15.8 s, total: 1min 58s
Wall time: 42.2 s

DataAugmentationしているのにも関わらず、このスピードが出てます。素晴らしい・・・

注意点

スピードの一番の要は、4Dイメージをバッチで一気に処理することによってオーバヘッドが減らせるっていうだけなので、
ImageDataGeneratorのようなバッチを排出するタイプだと遅いってことです。

おわりに

まとめると
- 一気にまとめると早いのでtf.data.Datasetはおすすめ
- だけど画像は荒いし、画像処理できるtf系はかなり限られている
- Augmentationの機能は今後に期待するしかない(自分が実装してTFAにプルリクしてもいいんですが...)

今後時間があれば、tf.data.Datasetに対応した画像処理系を作ろうかと思ってます。

追記

普通にこう書けばええんちゃう?

ImageDataGen
%%time
datagen = keras.preprocessing.image.ImageDataGenerator(
    horizontal_flip=True,
    vertical_flip=True,
    rotation_range=30,
)
for i,(x,y) in zip(range(5),datagen.flow(tr_x, tr_y, batch_size=tr_x.shape[0])):
    model.fit(x,y,batch_size=128,epochs=1)

確かに学習中は早いみたいです。

結果
Train on 50000 samples
50000/50000 [==============================] - 7s 132us/sample - loss: 0.9494 - accuracy: 0.6730
Train on 50000 samples
50000/50000 [==============================] - 7s 140us/sample - loss: 0.9275 - accuracy: 0.6797
Train on 50000 samples
50000/50000 [==============================] - 7s 132us/sample - loss: 0.9019 - accuracy: 0.6894
Train on 50000 samples
50000/50000 [==============================] - 7s 132us/sample - loss: 0.8842 - accuracy: 0.6968
Train on 50000 samples
50000/50000 [==============================] - 7s 132us/sample - loss: 0.8769 - accuracy: 0.6957
CPU times: user 1min 20s, sys: 4.96 s, total: 1min 25s
Wall time: 1min 24s

でも、Wall Time(実際にかかった時間)をみると遅いんですよこれが。
毎回Batchごとに負荷が1CPUにかかってしまっていて、ボトルネックになっている関係で重いです。
同じ処理でもtf.data.Datasetの方が早かったりします。

さらに追記

mapにkerasのpreprocessをちゃんぽんしたらKerasの贅沢な機能も使えるし早いんちゃうん?
という疑問が浮かんできました。早速実験します。

keras.preprocess.image
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras.preprocessing.image import random_rotation

labels = np.array([
    'airplane',
    'automobile',
    'bird',
    'cat',
    'deer',
    'dog',
    'frog',
    'horse',
    'ship',
    'truck'])

(tr_x,tr_y),(te_x,te_y)=keras.datasets.cifar10.load_data()
tr_x, te_x = tr_x/255.0, te_x/255.0
tr_y, te_y = tr_y.reshape(-1,1), te_y.reshape(-1,1)


def r_rotate(imgs, degree):
    pics=imgs.numpy()
    _=0
    if tf.rank(imgs)==4:
        X = np.zeros(pics.shape)
        for i,pic in enumerate(pics):
            X[i]=random_rotation(pic, degree, 0, 1, 2)

    elif tf.rank(imgs)==3:
        X=random_rotation(pics, degree, 0, 1, 2)
    return X,_
@tf.function
def random_rotate(imgs, label):
    x = tf.py_function(r_rotate,[imgs,30],[tf.float32,tf.int32])
    X = x[0]
    X.set_shape(imgs.shape)
    return X, label

tr_ds = tf.data.Dataset.from_tensor_slices((tr_x,tr_y)).shuffle(40000).batch(128).map(random_rotate)

plt.figure(figsize=(10,10),facecolor="white")
for b_img,b_label in tr_ds:
    for i, img,label in zip(range(25),b_img,b_label):
        plt.subplot(5,5,i+1)
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        plt.imshow(img)
        plt.xlabel(labels[label])
    break
plt.show()

CIFAR10-preprocess-image.png
うまくできているようです。
速度検証しましょう

speed_test
%%time
tr_ds = tf.data.Dataset.from_tensor_slices((tr_x,tr_y)).shuffle(40000)
tr_ds = tr_ds.batch(tr_x.shape[0]).map(random_rotate,num_parallel_calls=tf.data.experimental.AUTOTUNE).repeat(5)
tr_ds = tr_ds.prefetch(tf.data.experimental.AUTOTUNE)

for img,label in tr_ds:
    model.fit(x=img,y=label,batch_size=128)
結果
Train on 50000 samples
50000/50000 [==============================] - 9s 175us/sample - loss: 2.3420 - accuracy: 0.1197
Train on 50000 samples
50000/50000 [==============================] - 7s 131us/sample - loss: 2.0576 - accuracy: 0.2349
Train on 50000 samples
50000/50000 [==============================] - 7s 132us/sample - loss: 1.7687 - accuracy: 0.3435
Train on 50000 samples
50000/50000 [==============================] - 7s 132us/sample - loss: 1.5947 - accuracy: 0.4103
Train on 50000 samples
50000/50000 [==============================] - 7s 132us/sample - loss: 1.4540 - accuracy: 0.4705
CPU times: user 1min 33s, sys: 8.03 s, total: 1min 41s
Wall time: 1min 14s

前回ImageDataGeneratorで1min 30sくらいでしたから、若干高速に出来るくらいなようです。
ですが、自分のやり方が下手くそなので、分散処理できずにfor文で処理して1CPUに負荷がかかってしまっています。
なので、一旦はprefetchで高速化できたと行ったところでしょうか。
今後も実装していきます。

Suguru_Toyohara
渋谷でエンジニア。機械学習の仕事がメイン。でも他も結構できる。 興味ある技術:Vue.js/k8s/TF2.0/TFXあたり
microad
データとテクノロジーをかけ合わせたマーケティングプラットフォームを提供する会社です。
https://www.microad.co.jp/
Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
No comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
ユーザーは見つかりませんでした