最強データ拡張手法:RandAugmentを実装して、データ拡張の特性を調べてみよう

現在、最強のデータ拡張手法の一つRandAugmentをtf.data.Dataset向けに実装してみました。また、RandAugmentが、どのような挙動になるかを理解するために簡単な実験をしてみました。GWの結果を整理してみたいと思います。

RandAugmentとは？

画像データを回転や変形をしてデータを増やす、データをデータ拡張とよばれる手法があります。今回取り上げる、RandAugment: Practical automated data augmentation with a reduced search spaceは、Googleから提案されている最適なデータ拡張を探索する手法です。データ拡張の最適値を探す手法では、同じくGoogleから提案されたAutoAugmentなど手法があったのですが、最適な拡張となるパラメータを探索するコストが非常に大きいことが欠点でした。それを10の2乗オーダーで最適なデータ拡張を見つけ出すことを実現したのがRandAugmentです。NoisyStudentなどの論文でも採用されており、2020年5月現在、最強のデータ拡張手法の一つです。

というと、さぞかし難しそうに見えますが、RandAugumentは非常に単純なアルゴリズムです。RandAugmentは2つのパラメータ、n 、 mで制御されています。nはデータ拡張を何回行うかです。mはどれくらいの強さで拡張を行うかです。画像データを取り出すたびに、毎回、n個のデータ拡張操作をランダムで取り出し、mの大きさで加えます。

以下が論文からの図を参照してみます。１段目が大きさ9で2回拡張したもの、2段目が大きさ17、3段目が大きさ28で2回拡張させた画像です。

わずか2種類のパラメータで拡張の大きさを制御できるため、全探索しても、僅か10^2オーダーでおさまります。しかも探索できた最適値で行ったデータ拡張は、他手法によるデータ拡張を凌駕します。非常に簡単で強力な手法です。

RandAugumentを実装する

さて、RandAugmentを使った学習そのものは、2つのパラメータでグリッドサーチなので簡単に実装できます。

for n in range(拡張の種類):
   for m in range(拡張の最大値):
     rand_augment_training(n, m)

問題はデータ拡張を実施する側の方です。実質、実装と呼べるのは、こちらの部分です。
データ拡張を実施しているコードは、こちらにあるようです。

論文オリジナル実装@GitHub

なんとGoogleのコードなのにtf.data.Datasetで処理していないではないですか。。。まずはともあれ、こちらを参考にさせていただき、tf.data.Datasetで使用できるように変更をかけてみようと思います。tf.data.DatasetはTensorFlowのハイレベル APIの一つで、TensorFlowでの高速処理に向けては、今後はかかせなくなるデータ処理方法です。

参考：tf.data.Datasetのガイド

完全な再現からちょこっと目をつむった部分はありますが、自分の試作ライブラリのtftkに置いた通り、以下のように利用できるようになりました。とても簡単になりましたね。

n = 2  # 何種類の拡張をするか
m = 3  # どれくらいの大きさで

dataset:tf.data.Dataset  = tfds.load_dataset(xxxx) # 何かのtf.data.Dataset
dataset = dataset.map(ImageAugument.randaugment_map(n,m),num_parallel_calls=tf.data.experimental.AUTOTUNE)

実験:人力グリッドサーチでRandAugmentのお気持ちを理解する

さて、続いて実験です。Dogs vs. Catsのデータでやってみたいと思います。Dogs vs. Catsは画像を犬か猫かに分類するデータセットです。データは25000枚ありますが、そのうち、20000(犬猫1万ずつ)をトレーニングにまわしたいと思います。

参考：Dogs vs. Cats

ベンチマークは、こちらの記事にしてみましょうか。

上記のブログは、いろいろ模索しているので頑張り具合も伝わります。ファインチューニングを使わないと、以下あたりのAccuracyになるようですね。

Baseline VGG3 + Data Augmentation: 85.816

では、85.816%を目指していきましょう。

さて、RandAugmentのオリジナルな論文には重要な示唆が含まれています。一般に大きなデータセットに対して大きなデータ拡張が有効といったようなことが書かれています。Dogs vs. Catsの今回のトレーニングデータはわずか2万枚しかありません。ということは僅かな拡張にとどまる範囲でデータ拡張を考えてみた方がよさそうです。データ拡張は抑えめの範囲で考えてみましょう。

実装

まずは自作ライブラリですが、tftk(TensorFlow Tool kit)をインストールしてください。(いろいろ、最新物を突っ込もうとしますので、.venv環境をお使いください）

> pip install tftk -U

そして、Cats vs Dogsのデータを解凍しておいてください。
以下のコードでRandAugumentが実現できます。

import tensorflow as tf
import tftk
from tftk import Context
from tftk.image.dataset import ImageLabelFolderDataset
from tftk.image.dataset import ImageDatasetUtil
from tftk.image.augument import ImageAugument
from tftk.image.model import KerasResNet18
from tftk.image.model import SimpleClassificationModel
from tftk.train.image import ImageTrain
from tftk.callback import CallbackBuilder
from tftk.optimizer import OptimizerBuilder

if __name__ == '__main__':

    CLASS_NUM = 2
    IMAGE_SIZE = 150
    IMAGE_CHANNELS = 3
    EPOCHS = 100
    BATCH_SIZE = 100

    context = Context.init_context(TRAINING_NAME='DogsVsCats')
    train, train_len = ImageLabelFolderDataset.get_train_dataset(name="dogs-vs-cats", manual_dir="tmp")
    validation, validation_len = ImageLabelFolderDataset.get_validation_dataset(name="dogs-vs-cats", manual_dir="tmp")

    train = train.map(ImageDatasetUtil.map_max_square_crop_and_resize(IMAGE_SIZE,IMAGE_SIZE),num_parallel_calls=tf.data.experimental.AUTOTUNE)
    train = train.map(ImageAugument.randaugment_map(2,3),num_parallel_calls=tf.data.experimental.AUTOTUNE)
    train = train.map(ImageDatasetUtil.image_reguralization(),num_parallel_calls=tf.data.experimental.AUTOTUNE)
    train = train.map(ImageDatasetUtil.one_hot(CLASS_NUM),num_parallel_calls=tf.data.experimental.AUTOTUNE)

    validation = validation.map(ImageDatasetUtil.map_max_square_crop_and_resize(IMAGE_SIZE,IMAGE_SIZE),num_parallel_calls=tf.data.experimental.AUTOTUNE)
    validation = validation.map(ImageDatasetUtil.image_reguralization(),num_parallel_calls=tf.data.experimental.AUTOTUNE)
    validation = validation.map(ImageDatasetUtil.one_hot(CLASS_NUM),num_parallel_calls=tf.data.experimental.AUTOTUNE)

    optimizer = OptimizerBuilder.get_optimizer(name="rmsprop")
    model = KerasResNet18.get_model(input_shape=(IMAGE_SIZE,IMAGE_SIZE,IMAGE_CHANNELS),classes=CLASS_NUM)
    callbacks = CallbackBuilder.get_callbacks()
    ImageTrain.train_image_classification(
        train_data=train,train_size=train_len,
        batch_size=BATCH_SIZE,shuffle_size=100,
        validation_data=validation,validation_size=validation_len,
        model=model,callbacks=callbacks,
        optimizer=optimizer,loss="binary_crossentropy",max_epoch=EPOCHS)

それではResNet-18+RandAugmentで学習してみます。

まずはRandAugument(拡張の回数2、拡張の大きさ4)です。

Epoch 58/100
200/200 [==============================] - 33s 166ms/step - loss: 0.0524 - acc: 0.9140 - val_loss: 0.0510 - val_acc: 0.9312 - lr: 1.5625e-05

となりました。val_accは92.92%で、すでにベースラインを大きく越えてしまいました。では、これが最適なんでしょうか。

ただ、accが少し低いなぁと思います。データ拡張の手法なので複雑な画像が学習側に回っているので、学習側のaccが伸びづらいのは理解できます。しかし、どうもval_accもacc側の学習水準の低さに引っ張られている感もあるので、もう少し拡張を簡単にしてトレーニング側の方の数字も伸ばせば、val_accも伸びてくれないかな。。。ということで、拡張の大きさを少し落として学習を少し簡単にしてみましょう。

拡張の数を減らしてみます。ResNet(1-4)にしてみます。

Epoch 57/100
200/200 [==============================] - 25s 123ms/step - loss: 0.0329 - acc: 0.9719 - val_loss: 0.0530
 - val_acc: 0.9452 - lr: 1.5625e-05

おおっと、ベストがでましたね。85.816%からは大きく更新したように思います。

では、拡張の数の側を増やし、拡張の強さを弱めて、ResNet18 RandAug(3-2)してみました。
調子に乗りすぎました。まったく学習精度があがりませんでした。。。。拡張の数が増えるのは結構厳しそうです。

ついでに、nvidia-smiを使ってGPUの負荷を確認しておきましょう。こちらも95%でているので、きちんとGPUを使えていることがわかります。tf.data.Dataset#map()でそこそこ早く動いているのではないでしょうか。

>nvidia-smi
Sat May 02 01:02:01 2020
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 442.23       Driver Version: 442.23       CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce RTX 2060   WDDM  | 00000000:01:00.0 Off |                  N/A |
| N/A   83C    P2    78W /  N/A |   5330MiB /  6144MiB |     95%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     52128      C   ...cal\Programs\Python\Python37\python.exe N/A      |
+-----------------------------------------------------------------------------+

結局、Dogs vs Catsという簡単な小規模データセットでは、大きさ4で１度だけ拡張するがよかったようです。データが少ないだけに単純なところに落ち着いたのでしょうか。とはいえ、ベンチマークからは8%以上改善しております。なかなかの結果と言えそうですね。

まとめ

今回の実験で以下のようなことがわかりました。

• データ拡張の最強手法 RandAugmentをtf.data.Datasetで利用できるようにした。(tffkに入ってますよ)
• 少ないデータ量だとn, mは控えめにしよう。
• 拡張回数を増やすのは、拡張の強さを上げるより影響が大きそう。

自作ライブラリにつき、いろいろ改変されているかもしれませんが、機能追加、不具合、改善等、プルリク歓迎してます。