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

## 動作環境
* 今回は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に登録

```python
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）で学習させる

```python

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_trakking_mnist.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/23102/26321066-48da-8c2c-62b1-b86b47e7131b.png)
* MLflowのSearch APIでメトリクスのaccuracyが一番高いモデルのrun_idを取得
  * 実際にはCNNで学習したモデルのほうが高いaccuracyとなっている

```python

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に登録する

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

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

```python
from mlflow.tracking import MlflowClient
 
client = MlflowClient()
client.transition_model_version_stage(
  name=model_name,
  version=model_version.version,
  stage="Production"
)
```

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

```python
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]
```

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

```python
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_keras_test_labels.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/23102/949d09ad-5eec-e474-133b-6b4f96113ba4.png)

## 最後に
* 前回の[MLflow on Databricksを試してみた](https://qiita.com/sugattihide/items/f4b9a1b878cd43cc3885)ではDatabricksのNotebookからのMLflowでの運用方法を紹介しました。今回はAPIを利用した利用イメージを紹介してみました。