Keras(TensorFlow)のImageDataGeneratorをカスタマイズする(Mixup編)

こんちには。

データアナリティクス事業本部機械学習チームの中村です。

今回は、画像データの水増し(Data Augmentation)が可能な、KerasのImageDataGeneratorのカスタマイズ方法の第２弾です。

ImageDataGeneratorの基本的な使い方は、以下の前回記事を参照ください。

またカスタマイズの基礎となる第１弾は以下となっておりますので、併せてお読み頂ければと思います。

カスタマイズの概要

カスタマイズの概要は以下でした。

• ImageDataGeneratorを継承するクラスを作成する。
• flowという関数を修正する。
  • whileループ内でyieldで返す関数として実装する。
  • whileループ内ではカスタムでいれたい変換処理を実装する。

今回もこの方法でMixupという手法を導入してみます。 Mixupは、前回のRandom Erasingと異なり、複数の入力画像データを使ったデータ拡張手法となりますので、その点に着目してご覧ください。

Mixupについて

Mixupは2つの画像をλ:1-λの比率で混合した画像を生成するデータ拡張手法です。

以下のように画像を混合し、正解ラベルも同時に混合します。 ラベルの混合に対応するため、ラベルデータは混合前にOnehot Encodingして扱い、混合後は小数点の値を取る可能性があります。

混合率λは、ベータ分布B(α,α)をサンプリングして与えます。 ベータ分布の形状はαというパラメータで決まり、α=0.2の場合は以下のような形です。

横軸がλで、縦軸は度数になります。α=0.2の場合、ほとんどのケースでλ=0.0あるいは1.0に近いため、混合されにくい、混合されたとしても片方の割合がかなり小さいということになります。

α=1.0に変更すると以下のような一様分布となり、

α=1.5だと0.5に集まりますので、半々に混合される割合が多くなります。

Mixupを使う場合は、0.1～0.4の間で設定した方が良いと論文には記載されています。

このベータ分布のパラメータαがMixupのパラメータとなります。

詳細は以下の元論文を参照ください。

変換の実装

前準備

モジュールのインポートやデータ取得を行います。

import tensorflow as tf
import tensorflow.keras as keras
import numpy as np
import matplotlib.pyplot as plt

前回と同じく、データはCIFAR-10というRGB画像のデータセットを使います。

dataset = keras.datasets.cifar10
(train_images, train_labels), (test_images, test_labels) = dataset.load_data()

test_labelsをこの時点でOnehot Encodingしておきます。

train_labels = tf.one_hot(train_labels[:,0], 10) # train_labelsのshapeがなぜか(50000,1)だったので
train_labels.shape
# OUT: TensorShape([50000, 10])

この中から先頭16個のサンプルデータを持ってきておきます。

train_image_samples = train_images[:16]
train_label_samples = train_labels[:16]

カスタムクラス作成

以下のようにImageDataGeneratorを継承したカスタムクラスを作成します。

class MyImageDataGenerator(keras.preprocessing.image.ImageDataGenerator):

    def __init__(self,
        mixup_alpha=None,
        *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.mixup_alpha = mixup_alpha

    def mixup(self, X1, y1, X2, y2):
        # ...一旦省略...

    def flow(self, seed=None, *args, **kwargs):
        
        batch_gen = super().flow(seed=seed, *args, **kwargs)
        if seed is None:
            batch_gen2 = super().flow(seed=seed, *args, **kwargs)
        else:
            batch_gen2 = super().flow(seed=seed+777, *args, **kwargs) # seed+777はseedと同じでなければ何でも良い

        while True:

            batch_x, batch_y = next(batch_gen)
            batch_x_2, batch_y_2 = next(batch_gen2)

            if self.mixup_alpha is not None:
                batch_x, batch_y = self.mixup(batch_x, batch_y, batch_x_2, batch_y_2)

            yield (batch_x, batch_y)

__init__でカスタム処理で使うパラメータを引数に追加します。各パラメータは以下の通りです。

• mixup_alpha: ベータ分布のパラメータα

flowの部分は以下のようにしています。ここが複数の入力画像データを使うデータ拡張のポイントです。

• 継承元のflowを別々のseedで実行してgeneratorを２つ取得します。
• whileループ内では、それぞれのgeneratorからバッチデータを取得します。
• そして２種類のデータを使って混合処理をmixup関数で処理し、処理結果をyieldで返します。

Mixup処理は別途関数として以下のように定義します。

    def mixup(self, X1, y1, X2, y2):
        assert X1.shape[0] == y1.shape[0] == X2.shape[0] == y2.shape[0]
        batch_size = X1.shape[0]
        l = np.random.beta(self.mixup_alpha, self.mixup_alpha, batch_size)
        X_l = l.reshape(batch_size, 1, 1, 1)
        y_l = l.reshape(batch_size, 1)
        X = X1 * X_l + X2 * (1-X_l)
        y = y1 * y_l + y2 * (1-y_l)
        return X, y

描画用関数の準備

こちらは前回と似ていますが、画像は複製しなくなっています。

def plot_augmentation_image(train_image_samples, train_label_samples, params):

    # paramsで与えられた変換を実施
    data_generator = MyImageDataGenerator(**params)
    generator = data_generator.flow(
        x=train_image_samples, y=train_label_samples, batch_size=16)

    # 変換後のデータを取得
    batch_x, batch_y = generator.__next__()

    # 変換後はfloat32となっているため、uint8に変換
    batch_x = batch_x.astype(np.uint8)

    # 描画処理
    plt.figure(figsize=(10,10))
    for i in range(16):
        plt.subplot(4,4,i+1)
        plt.imshow(batch_x[i])
        plt.tick_params(labelbottom='off')
        plt.tick_params(labelleft='off')

描画して動作確認

まずは元論文でベースラインとして設定されているパラメータで検証します。

params = {
    'mixup_alpha': 0.2,
}
plot_augmentation_image(train_image_samples, train_label_samples, params)

正直まざっているか判断しずらいですが、最下段の右から二番目は何かMIXされていそうですね。

分かりやすいように、mixup_alphaを3.0当たりと大きめにしてみます。

params = {
    'mixup_alpha': 3.0,
}
plot_augmentation_image(train_image_samples, train_label_samples, params)

混ざっているものが増えました。正常に動いていそうです。

トレーニング

サンプルモデルで学習してみます。

params = {
    'mixup_alpha': 0.2,
}
data_generator = MyImageDataGenerator(**params)

batch_size = 32
generator = data_generator.flow(x=train_images, y=train_labels, batch_size=batch_size)

model = keras.Sequential([
    keras.layers.Conv2D(32, (3, 3), input_shape=(32, 32, 3), padding='same'),
    keras.layers.ReLU(),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Conv2D(32, (3, 3), padding='same'),
    keras.layers.ReLU(),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.GlobalAveragePooling2D(),
    keras.layers.Dense(10, activation='softmax'),
])

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

model.fit(generator, epochs=3, steps_per_epoch=len(train_images) // batch_size)

lossの部分はsparse_categorical_crossentropyではなく、categorical_crossentropyを使う必要があります。

これは、train_labelsをOnehotにした場合は、この設定にする必要があります。

また、前回同様、fitにはgeneratorを渡し、step_per_epochでエポックの終わりを明示的に指定します。

まとめ

いかがでしたでしょうか？今回は画像のデータ拡張手法のカスタマイズ第２弾ということで、複数の入力画像データを使用するデータ拡張手法について扱いました。これに似たケースを自身で実装する場合に本記事がお役に立てば幸いです。

複数サンプルを使用するものには、CutMix等他にもありますので、良ければそちらもトライしてみてください。