TensorFlow 2.0 時代の Keras API での画像分類器

TensorFlowを初期の頃から触っていて define-and-run の流儀にはそれなりに慣れてしまっていたけど、そろそろTensorFlowも2.0がreleaseされそうだし(2019.09時点で 2.0rc1) 新しいinterfaceも触っておかないと、と思って勉強してみた。

Effective TensorFlow 2.0 を読むと、major changesとして "Eager execution"、recommendationsとして"Keras layers and models"が紹介されている。 これからの時代はKeras APIを使ってEager executionでやっていく必要がありそうだ。

お題: 将棋駒画像の分類

昨年くらいから将棋の画像認識をやろうと思って 駒の画像データセットを作成 していた。今回はこれを使う。

各駒14種の先手・後手で28種、空白マスを加えて計29 classesを対象として、各classにつき約200〜300枚くらいずつ 96x96 のカラー画像 をラベル付きで用意してある。

datasetの準備

ラベル付きの画像たちを一定の割合で training, validation, test のdatasetに分割する。 今回は約 8:1:1 で分割して、

• training: 6277枚
• validation: 816枚
• test: 745枚

のdatasetが用意できた。

後述する tf.keras.preprocessing.image.ImageDataGenerator で使いやすいよう、各datasetを各label毎のディレクトリ以下に展開。

dataset
├── test
│   ├── BLANK
│   │   ├── 0de35ef1668e6396720e6fd6b22502b9.jpg
│   │   ├── 15767c0eb70db908f98a9ac9304227c8.jpg
│   │   ├── 18b0f691ac7b1ba6d71eac7ad32efdd5.jpg
│   │   ...
│   ├── B_FU
│   │   ├── 01aad8b7a32ca76e1ed82d72d9305510.jpg
│   │   ├── 04bc08425fb859f883802228e723c215.jpg
│   │   ├── 080950c64fb64840da3835d67fb969b8.jpg
│   │   ...
│   ├── B_GI
│   ...
├── training
│   ├── BLANK
│   ├── B_FU
│   ...
└── validation
    ├── BLANK
    ├── B_FU
    ...

transfer learning

まずは簡単に、学習済みの MobileNetV2 のモデルを使った転移学習をしてみる。

tf.keras.applications packageにはpre-trained modelが幾つか同梱されているので、それを呼び出すだけで簡単に使うことができる。

TensorFlow Hub にも同様に学習済みモデルが公開されていて再利用できるのだけど、現状では TensorFlow 2.0 向けのものは少なくて、 MobileNetV2 のものは input size が 224x224 に制限されているなどでちょっと使いづらい。 今後もっと整備されていくのかもしれないけど 今回は tf.keras.applications.MobileNetV2 を使うことにする。

INPUT_IMAGE_SIZE = (96, 96)

tf.keras.applications.MobileNetV2(
    input_shape=INPUT_IMAGE_SIZE + (3,),
    include_top=False,
    pooling='avg',
    weights='imagenet')

input_shapeのサイズは 96, 128, 160, 192, 224 のどれかで指定できて(defaultは224)、それに応じたimagenetでの学習済みモデルがダウンロードされて使われるようだ。

optionを指定しないと1000 classesの分類用logitsが出力されるが、転移学習にはそれは不要でその前段階の特徴量だけ抽出できれば良いので include_top=False を指定。 input_shape=(96, 96, 3) だと(None, 3, 3, 1280) の特徴量が出力されるようになる。 これをさらに pooling='avg' を指定することで平均値を取るpoolingにより (None, 1280) の2D Tensorが出力されるようになる(pooling='max'と指定することもできる)。 この1280個の値を特徴量ベクトルとして利用して結合層の部分だけ学習させて最適化していく。

feature extraction

base networkとなるMobileNetV2を固定させたまま使う場合は 入力画像に対する特徴量の出力は常に固定になるはずなので、先に全画像に対する特徴量を抽出しておくことが出来る。

Eager executionのおかげで単純に out = model(images).numpy() といった形で呼び出して出力の値を取得できる。

directoryを走査して画像データを読み込んでモデルへの入力値を作り、出力された特徴量ベクトルと対応するlabel indexをセットにして保存。どうせ後でshuffleして使うので順番は気にしない。 ここではnumpy arrayにしてnpzで保存する。

def dump_features(data_dir, features_dir):
    with open(os.path.join(data_dir, 'labels.txt'), 'r') as fp:
        labels = [line.strip() for line in fp.readlines()]

    model = tf.keras.applications.MobileNetV2(
        input_shape=INPUT_IMAGE_SIZE + (3,),
        include_top=False,
        pooling='avg',
        weights='imagenet')
    model.trainable = False

    features, label = [], []
    for root, dirs, files in os.walk(data_dir):
        for filename in files:
            image = tf.io.read_file(os.path.join(root, filename))
            image = tf.io.decode_image(image, channels=3)
            image = tf.image.convert_image_dtype(image, dtype=tf.float32)
            features.append(model(tf.expand_dims(image, axis=0)).numpy().flatten())
            label.append(labels.index(os.path.basename(root)))
    np.savez(os.path.join(features_dir, 'out.npz'), inputs=features, targets=label)

流石にCPUだとそこそこ時間がかかるが、数千件程度なら数分待てば完了する。

>>> import numpy as np
>>> npz = np.load('features/training.npz')
>>> npz['inputs']
array([[0.19352797, 0.        , 0.        , ..., 0.        , 1.0640727 ,
        2.7432559 ],
       [0.        , 0.        , 0.        , ..., 0.        , 1.4123839 ,
        3.057496  ],
       [0.18237096, 0.        , 0.04695423, ..., 0.33184642, 0.8117143 ,
        1.5288072 ],
       ...,
       [0.        , 0.        , 0.06801239, ..., 0.        , 0.13248181,
        1.6491733 ],
       [0.2550986 , 0.        , 0.10878149, ..., 0.        , 0.0785673 ,
        0.0311639 ],
       [0.        , 0.        , 0.        , ..., 0.        , 0.        ,
        2.0744066 ]], dtype=float32)
>>> npz['inputs'].shape
(6277, 1280)
>>> npz['targets']
array([24, 24, 24, ..., 14, 14, 14])
>>> npz['targets'].shape
(6277,)

あとはこの入出力に最適化するように結合層を学習させていけば良い。

tf.data.Dataset による学習データ準備

def dataset(category):
    npz = np.load(os.path.join(features_dir, f'{category}.npz'))
    inputs = npz['inputs']
    targets = npz['targets']
    size = inputs.shape[0]
    return tf.data.Dataset.from_tensor_slices((inputs, targets)).shuffle(size), size

training_data, training_size = dataset('training')
validation_data, validation_size = dataset('validation')

training用と validation用でそれぞれ tf.data.Dataset.from_tensor_slices に inputs と targets のtupleを渡すことで、Modelに与える訓練用入力データの準備ができる。

>>> for images, labels in training_data.batch(32).take(1):
...     print(images.shape, labels.shape)

(32, 1280) (32,)

tf.keras.Sequential によるModel定義

学習させるModelは tf.keras.Sequential に tf.keras.layers.Layer のlistを渡して記述していく。

with open(os.path.join(args.data_dir, 'labels.txt')) as fp:
    labels = [line.strip() for line in fp.readlines()]
classes = len(labels)

model = tf.keras.Sequential([
    tf.keras.layers.InputLayer((1280,)),
    tf.keras.layers.Dropout(rate=0.2),
    tf.keras.layers.Dense(
        classes,
        activation='softmax',
        kernel_regularizer=tf.keras.regularizers.l2(1e-4)),
])
model.summary()

入力となる1280の特徴量ベクトルに対しDropoutを入れつつDense layerで 29 classesの分類になるようにしているだけ。最終層のactivationはsoftmaxに。

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
dropout (Dropout)            (None, 1280)              0
_________________________________________________________________
dense (Dense)                (None, 29)                37149
=================================================================
Total params: 37,149
Trainable params: 37,149
Non-trainable params: 0
_________________________________________________________________

summaryでモデルの概要が簡単に分かって便利。

学習

まずはlossやmetricsの定義など。 Model.compile() で何を計測して何を減少させていくか、などを決定する。

model.compile(
    optimizer=tf.keras.optimizers.RMSprop(),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

この転移学習用に用意したdatasetではtargetsのlabelはindexを表す単一のint値なので、one-hot vectorに変換していない場合は tf.keras.losses.SparseCategoricalCrossentropy のように Sparseとついたものを使う。 one-hotな表現の値を使いたい場合はtf.keras.utils.to_categorical を利用して変換することも出来るようだ。 optimizerは RMSpropがデフォルトで使われるらしい。

ここまで準備できたら学習開始。 Model.fit() を呼ぶだけ。

history = model.fit(
    training_data.repeat().batch(batch_size),
    steps_per_epoch=training_size // batch_size,
    epochs=100,
    validation_data=validation_data.batch(batch_size),
    validation_steps=validation_size // batch_size,
    callbacks=[tf.keras.callbacks.TensorBoard()])
print(history.history)

与える学習データには上で作った training_data を使う。これは繰り返し使うので .repeat()を指定。 .batch() で処理するのでそれぞれ data_sizeをbatch_sizeで割ったsteps数を指定している。

callbacks で各epochが終わったときなどに特別な処理を挟むことが出来て、ここでは tf.keras.callbacks.TensorBoard() を渡すことで ./logs以下にTensorBoardで確認する用のlogデータを書き込んでくれるようになる。

Train for 98 steps, validate for 12 steps
Epoch 1/100
2019-09-16 00:13:08.160855: I tensorflow/core/profiler/lib/profiler_session.cc:184] Profiler session started.
98/98 [==============================] - 1s 14ms/step - loss: 2.2567 - sparse_categorical_accuracy: 0.3463 - val_loss: 1.4933 - val_sparse_categorical_accuracy: 0.5755
Epoch 2/100
98/98 [==============================] - 0s 3ms/step - loss: 1.2355 - sparse_categorical_accuracy: 0.6362 - val_loss: 1.0762 - val_sparse_categorical_accuracy: 0.6966
Epoch 3/100
98/98 [==============================] - 0s 3ms/step - loss: 0.8975 - sparse_categorical_accuracy: 0.7380 - val_loss: 0.8647 - val_sparse_categorical_accuracy: 0.7630
Epoch 4/100
98/98 [==============================] - 0s 3ms/step - loss: 0.7226 - sparse_categorical_accuracy: 0.7943 - val_loss: 0.7586 - val_sparse_categorical_accuracy: 0.7943
Epoch 5/100
98/98 [==============================] - 0s 3ms/step - loss: 0.6036 - sparse_categorical_accuracy: 0.8364 - val_loss: 0.6855 - val_sparse_categorical_accuracy: 0.8073
...

Dense Layerの部分だけの学習なのでとても速く、CPUでも1epochあたり0.3秒程度で あっという間に学習が進む。lossが減少し sparse_categorical_accuracyが増加していくのが見てとれる。

とりあえず 100 epoch回して TensorBoardで確認すると

loss:

sparse_categorical_accuracy:

橙がtraining, 青がvalidation。training_dataに対してはどんどん正答率が上がるが validation_dataに対しての結果は90%に届くか届かないか…程度のところで止まってしまうようだ。 まぁImageNetで学習したMobileNetV2が将棋駒画像に対してどれだけの特徴を捉えられているかというのを考えるとそんなものかな、という気はする。

学習後のModelを保存

学習したDense Layerを使って、base networkのMobileNetV2と繋げてまた新しく tf.keras.Sequential を作る。こうすることで、今度は (96, 96, 3)の入力に対して (29)の出力をする画像分類器として動くModelになる。

classifier = tf.keras.Sequential([
    tf.keras.applications.MobileNetV2(
        input_shape=INPUT_IMAGE_SIZE + (3,),
        include_top=False,
        pooling='avg',
        weights='imagenet'),
    model,
])
classifier.trainable = False
classifier.summary()

Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
mobilenetv2_1.00_96 (Model)  (None, 1280)              2257984
_________________________________________________________________
sequential (Sequential)      (None, 29)                37149
=================================================================
Total params: 2,295,133
Trainable params: 0
Non-trainable params: 2,295,133
_________________________________________________________________

簡単に繋ぎ直して使うことが出来て便利…。

あとはこのModelを丸ごと保存。.save() を呼ぶだけで良い。

classifier.save('transfer_classifier.h5')

test_dataでの評価

保存したModelを使って、学習には使っていない testのデータを使って精度を評価してみる。

feature extractionしたときと同様にディレクトリを走査し 画像ファイルとlabel indexをzipして tf.data.Datasetを生成する。

保存したModelは tf.keras.models.load_model() で読み込める。学習時と同様に loss に SparseCategoricalCrossentropy, metricsに SparseCategoricalAccuracy を指定してcompileして、評価で見るべき値を定める。

def evaluate(data_dir, model_path):
    with open(os.path.join(data_dir, 'labels.txt'), 'r') as fp:
        labels = [line.strip() for line in fp.readlines()]
    label_to_index = {label: index for index, label in enumerate(labels)}

    def load_image(image_path):
        image = tf.io.decode_jpeg(tf.io.read_file(image_path), channels=3)
        return tf.image.convert_image_dtype(image, tf.float32)

    image_paths = pathlib.Path(os.path.join(data_dir, 'test')).glob('*/*.jpg')
    image_paths = list(image_paths)
    label_index = [label_to_index[path.parent.name] for path in image_paths]
    images_ds = tf.data.Dataset.from_tensor_slices([str(path) for path in image_paths]).map(load_image)
    labels_ds = tf.data.Dataset.from_tensor_slices(label_index)
    test_data = tf.data.Dataset.zip((images_ds, labels_ds)).shuffle(len(image_paths))

    model = tf.keras.models.load_model(model_path)
    model.trainable = False
    model.compile(
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
    model.summary()

    test_result = model.evaluate(test_data.batch(1))
    print(test_result)

745/745 [==============================] - 13s 17ms/step - loss: 0.4405 - sparse_categorical_accuracy: 0.9114
[0.44054528336796983, 0.9114094]

今回のtransfer learningでのModelの、全745件のtest_data画像への正答率は 91.14% となることが分かった。

fine tuning

transfer learningでは90%前後の精度が限界のようだが、base networkのMobileNetV2部分も学習対象に含めて訓練していくとどうなるか。

Model定義

model = tf.keras.Sequential([
    tf.keras.applications.MobileNetV2(
        input_shape=INPUT_IMAGE_SIZE + (3,),
        include_top=False,
        pooling='avg',
        weights='imagenet'),
    tf.keras.layers.Dropout(rate=0.1),
    tf.keras.layers.Dense(
        len(labels),
        activation='softmax',
        kernel_regularizer=tf.keras.regularizers.l2(1e-4)),
])
model.summary()

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
mobilenetv2_1.00_96 (Model)  (None, 1280)              2257984
_________________________________________________________________
dropout (Dropout)            (None, 1280)              0
_________________________________________________________________
dense (Dense)                (None, 29)                37149
=================================================================
Total params: 2,295,133
Trainable params: 2,261,021
Non-trainable params: 34,112
_________________________________________________________________

MobileNetV2→Denseで繋げただけというModelの構造自体はで変わらない。 ただ今度はMobileNetV2の部分も学習対象とするので Trainable paramsが 37,149 がら 2,261,021 に激増している。

tf.keras.preprocessing.image.ImageDataGeneratorを使ってaugmentation

tf.keras.preprocessing というmoduleがあって、ここにはデータの前処理のためのutitityがある。

画像に対しても、以前は tf.image の様々なoperationを自分で組み合わせて作っていたdata augmentationの処理をまとめてやってくれるclassがある。

training_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    rotation_range=2,
    width_shift_range=2,
    height_shift_range=2,
    brightness_range=(0.8, 1.2),
    channel_shift_range=0.2,
    zoom_range=0.02,
    rescale=1./255)

このような形で、augmentationをかける際のパラメータ、レンジを指定して ImageDataGenerator を作る。 回転、縦横への移動、拡大縮小や明るさ変更など、色々指定できる。

このGeneratorに対して データを流し込んでいくことで iteratorが作られる。 既に展開してあるデータがあれば flow で、directoryだけ指定してそこから画像ファイルを読み取ってもらう場合は flow_from_directory で。

training_data = training_datagen.flow_from_directory(
    os.path.join(data_dir, 'training'),
    target_size=(96, 96),
    classes=labels,
    batch_size=batch_size)

こうすると、training_dataから読み出すたびに ImageDataGenerator生成時に指定したパラメータに従って変換をかけた画像たちがbatchで生成されて出力されるようになる。

極端にパラメータを大きくしてみると 以下のような感じで、傾いてたり白飛びしてたり様々。 このへんはどういう画像分類タスクを対象とするかによって適切なパラメータが異なるものになりそう。

ともあれ、この ImageDataGenerator からの出力を使って学習させていく。

model.compile(
    optimizer=tf.keras.optimizers.RMSprop(),
    loss=tf.keras.losses.CategoricalCrossentropy(),
    metrics=[tf.keras.metrics.CategoricalAccuracy()])

ImageDataGenerator.flow_from_directory は class_mode='categorical' がdefaultになっていて、これはtargetsの出力がone-hotな状態になる。ので、これを学習に使う場合は Sparse ではない CategoricalCrossentropy, CategoricalAccuracy を使うことになる。

そして学習。 training_data が ImageDataGenerator のものなので、Model.fit ではなく Model.fit_generator を使う。

history = model.fit_generator(
    training_data,
    epochs=100,
    validation_data=validation_data,
    callbacks=[
        tf.keras.callbacks.TensorBoard(),
        tf.keras.callbacks.ModelCheckpoint(
            os.path.join(weights_dir, 'finetuning_weights-{epoch:02d}.h5'),
            save_weights_only=True),
    ])

callbacks に tf.keras.callbacks.ModelCheckpoint() を追加してみた。epochごとにmodelの状態を保存してくれる。save_weights_only にすればmodel定義は無視して変数の値だけを保存してくれるようになる。

Epoch 1/100
99/99 [==============================] - 53s 539ms/step - loss: 0.9509 - categorical_accuracy: 0.7492 - val_loss: 4.0943 - val_categorical_accuracy: 0.3885
Epoch 2/100
99/99 [==============================] - 52s 520ms/step - loss: 0.2898 - categorical_accuracy: 0.9246 - val_loss: 8.5777 - val_categorical_accuracy: 0.1973
Epoch 3/100
99/99 [==============================] - 51s 515ms/step - loss: 0.2019 - categorical_accuracy: 0.9498 - val_loss: 7.5333 - val_categorical_accuracy: 0.2512
Epoch 4/100
99/99 [==============================] - 52s 521ms/step - loss: 0.1509 - categorical_accuracy: 0.9651 - val_loss: 11.6409 - val_categorical_accuracy: 0.0907
...

これは流石にCPUではかなりの時間がかかってしまう。手元のMacBookProだと 200s/epoch くらい。 Google Colaboratory の GPU Runtime を使うと 52s/epoch くらいに短縮できるようだ。 ちなみに TPU Runtime では TensorFlow 2.0をまだ使えなくて、そもそも2.0rc時点ではまだTPUをsupportできていないらしい。

ともあれ、どうにか 100 epoch回してみると

loss:

categorical_accuracy:

training(橙)は初期から順調にlossが減少してaccuracyも上昇するが、validation(青)はどうにも最初の数epochのうちは全然安定しない。 が、辛抱強く 30〜40epochくらいまで回していると突如としてlossの減少が始まり どんどん精度が上がってくる。 これはちょっとよく分からないけど fine-tuningでbase networkも学習するっていうのはこういうことなのかなぁ…。

ともかく、最終的にはかなり良い精度になっていそう。

評価

transfer learningのときと同じように model.save() で学習後のモデルを保存し、同じ評価scriptを使って evaluate を実行。

745/745 [==============================] - 14s 19ms/step - loss: 0.0542 - sparse_categorical_accuracy: 0.9906
[0.054189728634688225, 0.99060404]

転移学習より圧倒的に高い、 99.06% の精度が出た。すごい。 むしろ何を間違えているのかというと…

745件中 7件

(正解:△成香, 推論:△と金)

(正解:△香車, 推論:△金将)

(正解:△香車, 推論:△歩兵)

(正解:▲成香, 推論:▲成銀)

(正解:▲成香, 推論:▲と金)

(正解:▲成香, 推論:▲成銀)

(正解:▲歩兵, 推論: △成銀)

香車、成香はバリエーションあるわりにはデータあまり集めることができていなくて確かに間違うかもな、という感じ。とはいえ向き(先手か後手か)はだいたい見分けることが出来ている… と思ったら最後のやつは歩兵を後手の成銀と全然見当違いな結果になっていて謎。まぁ成銀もあまりデータ多くないので変なところで特徴を見てしまうのかも。

と考察できるくらいにはいいかんじに分類器として動いているようだ。

ここからは このModelを使ってJavaScriptやMobileAppで推論を動かしていく、というのをやっていくつもり

Repository

• https://github.com/sugyan/shogi-recognizer