tensorflow2のtf.dataを使ってaugmentationを高速化する

#はじめに
Kerasやtf.kerasのImageDataGeneratorは遅いので、`tf.data.Dataset`を使って学習を高速化してみます。
今回データ水増しには`Keras Preprocesing Layer`を使用します。**注：tensorflow2.3.0では使用可能ですが、まだ実験段階の機能とのことです。**なのでご注意ください。

環境
python 3.6.9
tensorflow 2.3.0
GPU　GTX1060

##参考文献
１．[TensorFlow公式チュートリアル](https://www.tensorflow.org/tutorials/load_data/images?hl=ja "TFチュートリアル")
チュートリアルらしく、step-by-stepでわかりやすいです。

２．[TensorFlowで使えるデータセット機能が強かった話](https://qiita.com/Suguru_Toyohara/items/820b0dad955ecd91c7f3)
tf.data.Datasetについてメチャクチャわかりやすい解説。とくに`shuffle`の説明がすごく良かったです。ありがとうございます。

３．[scikit-learn、Keras、TensorFlowによる実践機械学習 第2版](https://www.oreilly.co.jp/books/9784873119281/)
第２版になってメチャクチャ厚くなりました、、、。これも良いです。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/664963/dfb1df06-ee14-7278-d396-d1ec564bb250.png)

#練習用データセットの準備

公式サイトに準じて、５種類の花の画像のデータセットをDLします。

```py:python3
tf.keras.utils.get_file(origin='https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',fname='flower_photos', untar=True)

Downloading data from https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz
228818944/228813984 [==============================] - 23s 0us/step
'/home/username/.keras/datasets/flower_photos'
````

上記によって下記のディレクトリに画像がDLされます。
`/home/username/.keras/datasets/flower_photos`

今回は予めtrainingとvalidationに画像を分けてしまい、`training`というディレクトリに９割、`validation`というディレクトリに１割の画像が収まるように分けます。
まず`flower_photos`の作業中のディレクトリにコピーしてください。

```py:python3
import os
import shutil
import glob
import random

data_dir = './flower_photos'
directories = [i for i in glob.glob(os.path.join(data_dir, '*')) if os.path.isdir(i)]
```
`directories`には5つのディレクトリのパスが確認できます。
それぞれのディレクトリ名が画像のラベル（正解データ）になります。

```py:python3
def make_dir(path):
    if not os.path.exists(path):
        os.mkdir(path)

train_dir = './training'
valid_dir = './validation'
make_dir(train_dir)
make_dir(valid_dir)

for i in directories:
    files = glob.glob(os.path.join(i, '*jpg'))
    random.seed(42)
    random.shuffle(files)
    train_dir = os.path.join(train_dir, os.path.basename(i))
    valid_dir = os.path.join(train_valid, os.path.basename(i))
    make_dir(train_dir)
    make_dir(valid_dir)
    for k in files[int(len(files)*0.9):]:
        shutil.move(k, train_dir)
    for k in files[:int(len(files)*0.9)]:
        shutil.move(k, valid_dir)
```
これで各ラベルの9割がtraining、1割がvalidationのフォルダにわけられました。

諸々importします。

```py:python3
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import random
AUTOTUNE = tf.data.experimental.AUTOTUNE
```

画像を収めたディレクトリを確認していきます。

```py:python3
train_dir = './training'
valid_dir = './validation'
label_names = [os.path.basename(i) for i in glob.glob(os.path.join(train_dir, '*')) if os.path.isdir(i)]
label_names.sort()
print(label_names)

出力結果
['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
```
daisy, dandelion, roses, sunflowers, tulipsの５つのラベルに別れています。余談ですが、このデータセット、何故か花以外の画像が混入しています。。。

画像の枚数を確認します。

```py:python3
train_image_paths = glob.glob(os.path.join(train_dir, '*', '*jpg'))
valid_image_paths = glob.glob(os.path.join(valid_dir, '*', '*jpg'))
train_image_count = len(train_image_paths)
valid_image_count = len(valid_image_paths)
print('number of training image = ', train_image_count)
print('number of validation image = ', valid_image_count)

出力結果
number of training image =  3301
number of validation image =  369
```
Training画像が3301枚、Validation画像が369枚あります。


今度は、5つのラベルに任意のindex(番号)を割り付けます。どのように割り付けても良いと思います。

```py:python3
label_to_index = dict((name, index) for index,name in enumerate(label_names))
label_to_index

出力結果
{'daisy': 0, 'dandelion': 1, 'roses': 2, 'sunflowers': 3, 'tulips': 4}
```

すべての画像pathに対してラベル付けを行いリストとして格納します。

```py:python3
train_image_labels = [label_to_index[os.path.basename(os.path.dirname(path))]
                    for path in train_image_paths]
valid_image_labels = [label_to_index[os.path.basename(os.path.dirname(path))]
                    for path in valid_image_paths]

print("First 10 labels indices of train: ", train_image_labels[:10])
print("First 10 labels indices of validation: ", valid_image_labels[:10])

出力結果
First 10 labels indices of train:  [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
First 10 labels indices of validation:  [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
```
画像のpathはランダム化していないので、最初の10個をスライスすると、当然同じラベルになりました。

#####学習に向けてone-hot encodingする。
[`tf.one_hot`](https://www.tensorflow.org/api_docs/python/tf/one_hot)を使用して、one-hot encodingします。5クラスですので、引数のdepthは5にします。

```py:python3
train_image_labels = tf.one_hot(train_image_labels, depth=5)
valid_image_labels = tf.one_hot(valid_image_labels, depth=5)
print(train_image_labels[:10])

出力結果
tf.Tensor(
[[0. 1. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 1. 0. 0. 0.]], shape=(10, 5), dtype=float32)
```

#順番が前後しましたが、画像を確認してみます。


```py:python3
img_path = train_image_paths[3]
img_raw = tf.io.read_file(img_path)
img_tensor = tf.image.decode_image(img_raw)#uint8にデコードしてtensor化
plt.imshow(img_tensor)
```
出力結果
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/664963/8f5d2983-3568-b840-0ec1-361e6bb7d6b6.png)

＃画像の前処理関数を定義しておきます。
これはチュートリアルどおりです。

```py:python3
def preprocess_image(image):
    image = tf.image.decode_jpeg(image, channels=3)#channelsは1だとGrayscale、3だとRGB
    image = tf.image.resize(image, [150, 150])
    image /= 255.0  #機械学習のためnormalizationします。
    return image

def load_and_preprocess_image(path):
    image = tf.io.read_file(path)
    return preprocess_image(image)
```

#ようやくtf.data.Datasetを構築します。
さてここからが本題です。まず画像のpathを`from_tensor_slices`を使用して`tf.data.Dataset`

```py:python3
#まずtraining用画像のpathとlabelをそれぞれdataset化します。
path_ds = tf.data.Dataset.from_tensor_slices(train_image_paths)
label_ds = tf.data.Dataset.from_tensor_slices(train_image_labels)
#mapをつかってデータセット呼び出しに自動的に画像を変換するdatasetを作ります。
image_ds_train = path_ds.map(load_and_preprocess_image, num_parallel_calls=AUTOTUNE)
label_ds_train = tf.data.Dataset.from_tensor_slices(tf.cast(train_image_labels, tf.int64))
```

`Dataset`の挙動を確認がてら画像を確認します。

```py:python3
plt.figure(figsize=(8,8))
for n, image in enumerate(image_ds_train.take(10)):
    plt.subplot(5, 5, n+1)
    plt.imshow(image)
    plt.grid(False)
    plt.xticks([])
    plt.yticks([])
    plt.grid(False)  
plt.show()
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/664963/7e2fc8e2-9845-3ca5-c69e-0dac667fc04b.png)

`Dataset`はイテレータとして動作し、`image_ds_train`から`.take(10)`で10枚分の画像を呼び出し`load_and_preprocess_image`によって画像サイズをリサイズ、正規化したものが返されます。
なお`label_ds_train`はone-hot encodingされたラベルを返すイテレータとなります。

`image_ds_train`と`label_ds_train`は同じ順序ですので、`zip`することで(image, label) というペアのデータセットができます。便利ですねー。

```py:python3
image_label_ds_train = tf.data.Dataset.zip((image_ds_train, label_ds_train))
print(image_label_ds_train)

出力結果
<ZipDataset shapes: ((150, 150, 3), (5,)), types: (tf.float32, tf.int64)>
```



`image_label_ds_train`はイテレータとして使用すると各iterationごとに画像とラベルの2つを返します。
`shuffle()`については後述しますが、とりあえず挙動を確認してみます。

```py:python3
plt.figure(figsize=(15,10))
for n,image in enumerate(image_label_ds_train.shuffle(train_image_count).take(9)):
    plt.subplot(3, 3, n+1)
    plt.imshow(image[0])
    plt.grid(False)
    plt.xticks([])
    plt.yticks([])
    plt.xlabel(os.path.basename(str(image[1])))  
    plt.grid(False)
plt.show()
```
出力結果。one-hot encodingされたラベルと画像が一致していることが確認できました。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/664963/189dc755-1f60-fb0a-35c7-c9d3e714db5b.png)


さてパイプラインを完成させますが、ここで`shuffle()`について。
詳細は　[TensorFlowで使えるデータセット機能が強かった話](https://qiita.com/Suguru_Toyohara/items/820b0dad955ecd91c7f3)　を参考にしてください。
ポイントは
**1.`buffer_size`はデータセットと同じ数にすることで、データが完全にシャッフルされます。**
**2.``.repeat`の前に`.shuffle`すると1エポックの間に同じ画像が2回呼び出されるかもしれない。**引用：[TensorFlow公式チュートリアル](https://www.tensorflow.org/tutorials/load_data/images?hl=ja "TFチュートリアル")

まずtrainingのジェネレーターを完成させます。

```py:python3
BATCH_SIZE = 64

# シャッフルバッファのサイズをデータセットとおなじに設定することで、データが完全にシャッフルされる
# ようにできます。
ds = image_label_ds_train
ds = ds.repeat()
ds = ds.shuffle(buffer_size=train_image_count) #shuffleの順序注意！
ds = ds.batch(BATCH_SIZE)
ds = ds.prefetch(buffer_size=AUTOTUNE)
ds
```

validationも同様に作成しますが、validationはshuffleする必要はないので、`shuffle`なし。

```py:python3
path_ds_valid = tf.data.Dataset.from_tensor_slices(valid_image_paths)
image_ds_valid = path_ds_valid.map(load_and_preprocess_image, num_parallel_calls=AUTOTUNE)
label_ds_valid = tf.data.Dataset.from_tensor_slices(tf.cast(valid_image_labels, tf.int64))
image_label_ds_valid = tf.data.Dataset.zip((image_ds_valid, label_ds_valid))

ds_valid = image_label_ds_valid
ds_valid = ds_valid.repeat()
ds_valid = ds_valid.batch(BATCH_SIZE)
ds_valid = ds_valid.prefetch(buffer_size=AUTOTUNE)
ds_valid
```

#学習してみる。

適当にモデルを作成してコンパイルします。

```py:python3
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPooling2D,BatchNormalization
from tensorflow.keras.preprocessing.image import ImageDataGenerator
IMG_HEIGHT = 150
IMG_WIDTH = 150
epochs = 20
model = Sequential([
    Conv2D(16, 3, padding='same', activation='relu', input_shape=(IMG_HEIGHT, IMG_WIDTH ,3)),
    BatchNormalization(),
    MaxPooling2D(),
    Conv2D(32, 3, padding='same', activation='relu'),
    BatchNormalization(),
    MaxPooling2D(),
    Conv2D(64, 3, padding='same', activation='relu'),
    BatchNormalization(),
    MaxPooling2D(),
    Dropout(0.5),
    Flatten(),
    Dense(512, activation='relu'),
    Dense(5, activation='softmax')
])

model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])
```

`model.fit`で学習します。kerasでよくお世話になった`fit_generator`はもう使わないようです。
`max_queue_size=120, workers=30, use_multiprocessing=True,`の部分はお使いのcpuやosに合わせてください。

```py:python3
history = model.fit(ds, 
        epochs=epochs, 
        steps_per_epoch=train_image_count//BATCH_SIZE,
        validation_data=ds_valid,
        validation_steps=valid_image_count//BATCH_SIZE, 
        validation_batch_size=BATCH_SIZE,
        max_queue_size=120, workers=30, use_multiprocessing=True,
        )
```

学習が始まります。

```
Epoch 1/20
 1/51 [..............................] - ETA: 0s - loss: 2.9178 - accuracy: 0.2656WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0157s vs `on_train_batch_end` time: 0.0288s). Check your callbacks.
51/51 [==============================] - 2s 47ms/step - loss: 3.2386 - accuracy: 0.4960 - val_loss: 4.5163 - val_accuracy: 0.3781
・
・
・
Epoch 20/20
51/51 [==============================] - 2s 45ms/step - loss: 0.0474 - accuracy: 0.9835 - val_loss: 1.6496 - val_accuracy: 0.6438
```

結果をplotします。

```py:python3
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
 
loss = history.history['loss']
val_loss = history.history['val_loss']
 
epochs_range = range(epochs)
 
plt.figure(figsize=(16, 8))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')
 
plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label='Training Loss')
plt.plot(epochs_range, val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()
```

結果。過学習しているのか、training と　validationの曲線に乖離が見られました。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/664963/35550768-547c-6ea3-19d0-2bc7804f9067.png)

#tf.data.Datasetにaugmentationを組み合わる。

TensorFlowの公式チュートリアルを参考にしました。
[データの増強](https://www.tensorflow.org/tutorials/images/data_augmentation?hl=ja)
`Keras Preprocesing Layer`を使用します。
**注：tensorflow2.3.0では使用可能ですが、まだ実験段階の機能とのことです。**

まず前処理レイヤーを作成します。

```py:python3
data_augmentation = tf.keras.Sequential([
  tf.keras.layers.experimental.preprocessing.RandomFlip("horizontal_and_vertical"),
  tf.keras.layers.experimental.preprocessing.RandomRotation(0.2,fill_mode = 'reflect'),
  tf.keras.layers.experimental.preprocessing.RandomZoom(height_factor=0.2, fill_mode =  'reflect'),
])
```
1枚画像を使用して、水増しできているか確認します。

```py:python3
#学習画像を1枚読み込みます。
img_path = train_image_paths[3]
img_raw = tf.io.read_file(img_path)
img_tensor = tf.image.decode_image(img_raw)#uint8にデコードしてtensor化
plt.imshow(img_tensor)
```
元画像
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/664963/128a2cc8-ee78-4b11-046a-19bca3fcf50f.png)

```py:python3
img = tf.expand_dims(img_tensor, 0)
plt.figure(figsize=(10, 10))
for i in range(9):
    augmented_image = data_augmentation(img)
    ax = plt.subplot(3, 3, i + 1)
    plt.imshow(np.asarray(augmented_image[0], dtype=np.uint8))
    plt.axis("off")
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/664963/e7d9570b-056e-3646-e219-d52f50634a84.png)


どうやら水増しできているようです。
水増し手法は結構網羅されているようで、[ここ](https://www.tensorflow.org/api_docs/python/tf/keras/layers/experimental/preprocessing?hl=ja)から確認できます。
水増ししたジェネレーターを作成します。

```py:python3

ds_aug = image_label_ds_train
ds_aug = ds_aug.repeat()
ds_aug = ds_aug.shuffle(buffer_size=train_image_count) #shuffleの順序注意！
ds_aug = ds_aug.batch(BATCH_SIZE)
ds_aug = ds_aug.map(lambda x, y: (data_augmentation(x, training=True), y))
ds_aug = ds_aug.prefetch(buffer_size=AUTOTUNE)
```
モデルは先程と同じものを使用して学習しますが、重みを再利用してしまうので、新しく定義してコンパイルします。

```py:python3
model_aug = Sequential([
    Conv2D(16, 3, padding='same', activation='relu', input_shape=(IMG_HEIGHT, IMG_WIDTH ,3)),
    BatchNormalization(),
    MaxPooling2D(),
    Conv2D(32, 3, padding='same', activation='relu'),
    BatchNormalization(),
    MaxPooling2D(),
    Conv2D(64, 3, padding='same', activation='relu'),
    BatchNormalization(),
    MaxPooling2D(),
    Dropout(0.5),
    Flatten(),
    Dense(512, activation='relu'),
    Dense(5, activation='softmax')
])

model_aug.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

history = model_aug.fit(ds_aug, 
        epochs=epochs, 
        steps_per_epoch=train_image_count//BATCH_SIZE,
        validation_data=ds_valid,
        validation_steps=valid_image_count//BATCH_SIZE, 
        validation_batch_size=BATCH_SIZE,
        max_queue_size=120, workers=30, use_multiprocessing=True,
        )
```
学習結果

```
Epoch 1/20
51/51 [==============================] - 3s 55ms/step - loss: 3.3589 - accuracy: 0.4887 - val_loss: 11.0810 - val_accuracy: 0.2812
Epoch 2/20
51/51 [==============================] - 2s 46ms/step - loss: 1.1092 - accuracy: 0.5665 - val_loss: 16.0235 - val_accuracy: 0.2812
・
・
・
Epoch 20/20
51/51 [==============================] - 2s 47ms/step - loss: 0.7122 - accuracy: 0.7172 - val_loss: 0.6610 - val_accuracy: 0.7719
```

先ほどと同様に学習曲線をplotしてみると乖離が解消されており、データ水増しの効果が確認できました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/664963/cdb42a8f-7b71-8f48-fe1d-835c786b62f6.png)

###tf.data.Dataset＋Keras Preprocesing Layerとtf.kerasのflow+ImageDataGeneratorの速度比較

さて肝心の高速化はできてるのでしょうか？
同じ条件で旧来のtf.keras flow +  ImageDataGeratorで学習を実行してみたところ、
**1stepの速度**
`tf.data.Dataset` : 46 - 48 msec
`ImageDataGerator + flow` : 65 - 69 msec
**1epochの速度**
`tf.data.Dataset` : 2 - 3 sec
`ImageDataGerator + flow` : 3 - 4 sec

と僅かながら高速化されています。僅かですが、もっと多量の画像を扱う場合はチリツモでしょう。
加えて`tf.data.Dataset`を使用した場合は、epoch間も早く全体的にはかなり高速化されているように感じます。

#おわりに
まだ実験段階とのことですが、`Keras Preprocesing Layer`を使用して無事に水増しができました。はやく正式実装されるといいのですが・・・。`tf.data.Dataset`はとっつきにくいですが、なれれば`flow`系の手法より高速に学習できますので、おすすめです！