kerasを使ったMuti-task Learning(CNN + Autoencoder)

最新のモデルでは一般的になってきているMuti-taskなモデルについて取り上げたいと思います。

Multi-task Learningとは

Muti-task Learning(MTL)とは、１つのネットワークで複数のタスクを解くモデルです。
MTLの内容はこちらの記事にわかりやすくまとめられています。
直感的な理解としては、人は新しいタスクを学習しようとするとき、これまでの知識を活用することでスムーズに新しいタスクを学習することができます。例えば、将棋ができる人とできない人が新しくチェスを覚えようと思ったとき、将棋ができる人の方が覚えが早いことが多いと思います。
このように別の分野の知識を複合的に学習することで、学習のスピードや最終的な精度を上げようということがMTLのモチベーションです。

今回のモデル

CNNによるクラス分類とAutoencoderを同時に解くモデルを紹介します。
データセットにはMNISTを使います。
ソースコードはこちらです。

環境

• Python 3.6.4
• Keras 2.1.2
• tensorflow 1.4.1 (backend)

モデル構造

CNNの部分は前回の記事で扱った構造と同じです。
６層の畳み込みで特徴量抽出をし、MLPによるクラス分類とUpsamplingとConvolutionによるAutoencoderの２つに分岐させています。

model.py

def prepare_model(input_shape=(28, 28, 1), class_num=10):
    input = Input(input_shape)
    kernel_size = (3, 3)
    max_pool_size = (2, 2)
    upsampling_size = (2, 2)

    enc_cnn = Conv2D(64, kernel_size, padding='same', activation='relu')(input)
    enc_cnn = Dropout(0.1)(enc_cnn)
    enc_cnn = Conv2D(64, kernel_size, padding='same', activation='relu')(enc_cnn)
    enc_cnn = Dropout(0.1)(enc_cnn)
    enc_cnn = Conv2D(64, kernel_size, padding='same', activation='relu')(enc_cnn)
    enc_cnn = Dropout(0.1)(enc_cnn)
    enc_cnn = MaxPooling2D(pool_size=max_pool_size, strides=(2, 2))(enc_cnn)
    
    enc_cnn = Conv2D(64, kernel_size, padding='same', activation='relu')(enc_cnn)
    enc_cnn = Dropout(0.1)(enc_cnn)
    enc_cnn = Conv2D(64, kernel_size, padding='same', activation='relu')(enc_cnn)
    enc_cnn = Dropout(0.1)(enc_cnn)
    enc_cnn = Conv2D(64, kernel_size, padding='same', activation='relu')(enc_cnn)
    enc_cnn = MaxPooling2D(pool_size=max_pool_size, strides=(2, 2))(enc_cnn)

    fc = Flatten()(enc_cnn)
    fc = Dense(1024, activation='relu')(fc)
    softmax = Dense(class_num, activation='softmax', name='classification')(fc)

    dec_cnn = UpSampling2D(upsampling_size)(enc_cnn)
    dec_cnn = Conv2D(64, kernel_size, padding='same', activation='relu')(dec_cnn)
    dec_cnn = UpSampling2D(upsampling_size)(dec_cnn)
    dec_cnn = Conv2D(1, kernel_size, padding='same', activation='sigmoid', name='autoencoder')(dec_cnn)

    outputs = [softmax, dec_cnn]

    model = Model(input=input, output=outputs)
    return model

一般的にクラス分類と画像復元(Autoencoder)では、損失関数にcross entropyとmean square error(二乗誤差)を用いるので損失関数をそれぞれ定義する必要があります。
Kerasでは出力層で定義した層に対する損失関数を辞書型で定義することができます。

train.py

    model.compile(loss={'classification': categorical_crossentropy, 'autoencoder': mean_squared_error},
                  loss_weights={'classification': 0.9, 'autoencoder': 0.1},
                  optimizer=SGD(lr=0.01, momentum=0.9, nesterov=True))

また、画像復元の方が学習の進みが早かったのでloss_weightsのパラメータを用いてクラス分類と画像復元を9:1の割合で学習させています。

学習結果

20エポック学習させた結果、テストデータに対する認識率は99.33%となりました。
また、復元した画像は以下のようになりました。
左側が元画像で右側が復元した画像になります。

まとめ

基本的なソースコードの実行の方法は前回とほとんど同じです。
今回は出力層で分岐させ複数のタスクを解くようなモデルを紹介しました。
少し複雑なモデルを組もうとした場合に、入力や出力を複数に分岐させるような場合がでてくると思うので今回のモデルのサンプルコードが参考になればと思います。