Kerasで作成したモデルをMLflowで管理してみた #Keras

動作環境

今回はDatabricks環境上で検証しています。記事作成時点（2020年11月）で最新RuntimeのDatabricks 7.4MLを利用しています
7.4MLクラスタの以下ライブラリを利用しています（初期インストール済）
- MLflow 1.11.0
- TensorFlow 2.3.1
- Keras 2.4.0

この記事の内容

以下をMLflowでトラッキングする
- MNISTをシンプルな３層レイヤーで学習させる
- MNISTをCNNで学習させる
最もAccuracyが高いモデルをModelレジストリに登録する
ModelをProductionステージへ変更する
モデルをロードして推論を実行する

MNISTをシンプルな3層レイヤーで学習させる

入力層に28×28pixelの手書き文字画像を入力
中間層に512個のノードを定義し、活性化関数はreluを設定
出力は10種類（0から9まで）に分類する
モデルパラメータ、メトリックをMLflowに登録

import tensorflow as tf
from tensorflow.keras import datasets, layers, models
import mlflow

(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()

train_images = train_images.reshape((60000, 28, 28, 1))
test_images = test_images.reshape((10000, 28, 28, 1))

# ピクセルの値を 0~1 の間に正規化
train_images, test_images = train_images / 255.0, test_images / 255.0

with mlflow.start_run(run_name='keras-mnist'):
  model = models.Sequential()
  model.add(layers.Flatten())
  model.add(layers.Dense(512, activation='relu'))
  model.add(layers.Dropout(0.2))
  model.add(layers.Dense(10, activation='softmax'))

  model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

  batch_size = 128
  epoch = 5

  model.fit(train_images, train_labels,
          batch_size=batch_size,
          epochs=epoch)

  mlflow.log_param("batch_size", batch_size)
  mlflow.log_param("epoch_num", epoch)

  test_loss, test_acc = model.evaluate(test_images,  test_labels, verbose=2)
  mlflow.log_metrics({'loss': test_loss})
  mlflow.log_metrics({'accuracy': test_acc})
  mlflow.keras.log_model(model, "keras-mnist")

MNISTをCNNで学習させる

同じMNISTの分類を畳み込みニューラルネットワーク（CNN）で学習させる


import tensorflow as tf
from tensorflow.keras import datasets, layers, models
import mlflow

(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()

train_images = train_images.reshape((60000, 28, 28, 1))
test_images = test_images.reshape((10000, 28, 28, 1))

train_images, test_images = train_images / 255.0, test_images / 255.0

with mlflow.start_run(run_name='keras-mnist'):
  model = models.Sequential()
  model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)))
  model.add(layers.MaxPooling2D((2, 2)))
  model.add(layers.Conv2D(64, (3, 3), activation='relu'))
  model.add(layers.MaxPooling2D((2, 2)))
  model.add(layers.Conv2D(64, (3, 3), activation='relu'))
  model.add(layers.Flatten())
  model.add(layers.Dense(64, activation='relu'))
  model.add(layers.Dense(10, activation='softmax'))

  epoch = 5

  model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
  model.fit(train_images, train_labels, epochs=epoch)

  mlflow.log_param("epoch_num", epoch)

  test_loss, test_acc = model.evaluate(test_images,  test_labels, verbose=2)
  mlflow.log_metrics({'loss': test_loss})
  mlflow.log_metrics({'accuracy': test_acc})
  mlflow.keras.log_model(model, "keras-mnist")

最もAccuracyが高いモデルをModelレジストリに登録する

ここまででMLflowトラッキングによりモデルの実験結果が以下のように管理されている
MLflowのSearch APIでメトリクスのaccuracyが一番高いモデルのrun_idを取得
- 実際にはCNNで学習したモデルのほうが高いaccuracyとなっている


from mlflow.entities import ViewType

run_id = mlflow.search_runs(
  filter_string='tags.mlflow.runName = "keras-mnist"', 
  run_view_type=ViewType.ACTIVE_ONLY,
  max_results=1, 
  order_by=["metrics.accuracy DESC"]).iloc[0].run_id

取得したrun_idに紐づくモデルをModel Registryに登録する

model_name = "keras-mnist"
model_version = mlflow.register_model(f"runs:/{run_id}/keras-mnist", model_name)

ModelをProductionステージへ変更する

以下のクライアントAPIからモデルをProductionステージに変更させる
- クライアント実行ユーザーの権限でモデルのステージ変更が承認される

from mlflow.tracking import MlflowClient

client = MlflowClient()
client.transition_model_version_stage(
  name=model_name,
  version=model_version.version,
  stage="Production"
)

モデルをロードして推論を実行する

Model RegistryからProductionステージのモデルを取得して推論を実行する

import mlflow
from tensorflow.keras import datasets

model_name = "keras-mnist"
model = mlflow.keras.load_model(f"models:/{model_name}/production")

(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()
test_images = test_images.reshape((-1, 28, 28, 1))

result = model.predict(test_images[0:10])
print(result.argmax(axis=1))

テストデータで推論させると以下となっている

[7 2 1 0 4 1 4 9 5 9]

テストデータのラベルを可視化して確認してみると推論結果と一致した

import matplotlib.pyplot as plt
for i in range(10):
    plt.subplot(1, 10, i+1)
    plt.imshow(test_images[i].reshape(28,28))

最後に

前回のMLflow on Databricksを試してみたではDatabricksのNotebookからのMLflowでの運用方法を紹介しました。今回はAPIを利用した利用イメージを紹介してみました。