KerasでQuickDrawデータセットを使って画像分類問題を解く

はじめに

この記事では、TensorFlowで提供されているQuickDrawデータセットを使用し、学習したモデルを作成する方法をご紹介します。作成した学習モデルは最終的にVR機器であるMetaQuestに導入し、VRアプリで使用します。そのため、作成モデルはKerasモデルではなく、ONNX形式で保存します。

学習環境

Windows 10(試作環境はWindows 11)
Python 3.10.12
TensorFlow 2.17.0

学習環境の構築

機械学習の環境構築は要件や手持ちの計算機スペック、金銭面などで大きく変わります。学習環境の構築は、以下の手順を参考にすることをお勧めします。

長時間計算機を使えない環境や計算機スペックが低く膨大な時間が必要になる場合 => クラウドサービスを利用する
Windows10でWSL2を使えない場合 => Linux環境をデュアルブートなどで構築する
WindowsでWSL2が利用可能な場合 => WSL2で構築する
そこまで計算機スペックが要求されないが、ローカルで構築したくなく、クラウドサービスを利用したくない場合 => Raspberry Piなどの小型計算機を利用する

今回は、WSL2を利用して環境構築を行います。手順は以下の通りです。

WSL2にUbuntu22.04をインストール
GPUを利用するのでCudaToolkitをインストール(今回は12.6を入れます)
Cudnnをインストール(今回は9.5.0を入れます)
TensorFlow(GPU)をインストール

利用するデータセット

QuickDrawデータセットを使用しますが、画像分類のモデルを作るため、QuickDrawBitmapデータセットを利用します。Tensorflowで読み込むには下記のようにします。

(train_ds, val_ds, test_ds), metadata = tfds.load(
    'quickdraw_bitmap',
    split=['train[:80%]', 'train[80%:90%]', 'train[90%:]'],
    with_info=True,
    as_supervised=True,
    data_dir='./datasets',
)

学習用を8割、検証用1割、テスト用に1割に分割しています。データセットのラベル数は執筆時点で345種類になります。

データセットの前処理

今回はEfficientNetV2B2を利用するため、画像サイズを(32, 32)にリサイズします。バッチサイズは128とします。

resize = tf.keras.Sequential([
    layers.Resizing(32, 32),
])

def normalize_img(image, label):
    _image = resize(image)
    return tf.cast(_image, tf.float32) / 255., label

train_ds = train_ds.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE).batch(128).prefetch(tf.data.AUTOTUNE)
val_ds = val_ds.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE).batch(128).prefetch(tf.data.AUTOTUNE)
test_ds = test_ds.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE).batch(128).prefetch(tf.data.AUTOTUNE)

モデルの作成/学習

前提としてMetaQuestの端末で動かすことを想定しているため、モデル容量などを考慮する必要があり、軽量なEfficientNetでニューラルネットを構築します。Kerasを利用するとEfficientNetV2B2を簡単に利用できます。

model = tf.keras.applications.EfficientNetV2B2(
    input_shape=(
        32,
        32,
        1,
    ),
    weights=None,
    classes=345,
)
model.compile(
    optimizer="adam",
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    metrics=["accuracy"],
)

チェックポイントは下記のように設定し、学習を行います。

cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1)
epochs = 50

history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=epochs,
    callbacks=[cp_callback]
)

モデルの評価

学習したモデルの性能を評価します。

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=(8, 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.savefig(output_result_img_location)

学習モデルの保存

ONNX形式でモデルを作りたいため、一度TensorFlow SavedModel形式で保存します。

tf.saved_model.save(model, './my_saved_model')

ONNX形式へ変換するために、tf2onnxを利用します。

python3 -m tf2onnx.convert --saved-model my_saved_model --output my_saved_model.onnx --opset 18 --verbose

最後に

今回学習したモデルは、MetaQuestで動かし、数人に体験してもらった結果大体6割程度の正解率でした。現在では、機械学習で生成するよりも生成AIを利用する流れが主流ですが、たまには自分で学習モデルを作成し、デバイスに組み込んで遊ぶのは楽しいものなので、ぜひ試してみてください。