MLflow を使った機械学習のライフサイクル管理

MLflowとは、機械学習のライフサイクルを管理するためのオープンソースのプラットフォームです。MLflow は、データサイエンティストやエンジニアが機械学習のプロジェクトを効率的に実行できるように設計されています。

MLflow の主な機能は以下のようなものがあります。

MLflow は上記のような機能により、機械学習プロジェクトの管理、実行、デプロイメントの効率を大幅に向上させることができます。

この記事では MLflow の基本的な使い方を説明します。

MLflow の環境準備

MLflow の環境準備方法ついて説明します。

MLflow のインストール

MLflow を使用する場合にはインストールが必要です。以下のように pip でインストールしてください。

pip install mlflow

MLflow サーバーの実行

MLflow サーバーを実行して利用します。リモートサーバーを立てる方法やローカルで実行する方法がありますが、この記事では MLflow の使い方を紹介するのが主目的であるため、ローカルで実行する方法を紹介します。

実行する python プログラムがあるフォルダで以下のコマンドを実行してください。

mlflow ui 

【実行結果】
>mlflow ui
INFO:waitress:Serving on http://127.0.0.1:5000

実行するとサーバーのURLが表示されます。http://127.0.0.1:5000 または http://localhost:5000/ で MLflow UI にアクセスできます。

MLflow UI は、機械学習の実行結果を Web ブラウザ上で管理するインターフェースです。これにより、異なる実験結果を視覚的に比較し、効果的な分析が可能になります。

MLflow の使用方法

MLflow の使用方法を簡単な例を使って紹介します。今回は、MNIST という手書き画像データ分類を例にします。なお、分類手法は CNN（Convolutional Neural Network：畳み込みニューラルネットワーク）で、実装フレームワークとしては Tensorflow / Keras を使用します。

CNN に関する説明はこの記事ではしません。CNN による画像分類の説明は「CNN（畳み込みニューラルネットワーク）による画像分類の基本」を参考にしてください。

MLflow を使用した実装例の全体は以下のようになります。ポイントとなる部分の詳細を以降で説明していきます。

import numpy as np
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.datasets import mnist

import mlflow
import mlflow.keras


def main():
    """メイン関数"""
    # ===== MLflowの実験設定
    mlflow.set_experiment("mnist_cnn_classification")

    # ===== MNIST(エムニスト)データの読込
    (train_imgs, train_labels), (test_imgs, test_labels) = mnist.load_data()
    train_imgs = train_imgs.reshape((60000, 28, 28, 1))
    test_imgs = test_imgs.reshape((10000, 28, 28, 1))
    # 訓練データの一部(20%)を評価データとして使う
    idx = int(train_imgs.shape[0] * 0.2)
    train_imgs, val_imgs = train_imgs[idx:], train_imgs[:idx]
    train_labels, val_labels = train_labels[idx:], train_labels[:idx]

    # ===== CNNモデルの構築
    # MNIST画像は28×28でチャンネルは1
    inputs = keras.Input(shape=(28, 28, 1))
    # 前処理0~1へ正規化
    x = layers.Rescaling(1.0 / 255)(inputs)
    # 畳み込み層とプーリング層の定義
    x = layers.Conv2D(32, kernel_size=3, activation="relu")(x)
    x = layers.MaxPooling2D(pool_size=2)(x)
    x = layers.Conv2D(64, kernel_size=3, activation="relu")(x)
    x = layers.MaxPooling2D(pool_size=2)(x)
    x = layers.Conv2D(128, kernel_size=3, activation="relu")(x)
    x = layers.MaxPooling2D(pool_size=2)(x)
    # 平坦化する
    x = layers.Flatten()(x)
    # ドロップアウトを設定
    x = layers.Dropout(0.5)(x)
    # 分類のために10のノードに接続
    outputs = layers.Dense(10, activation="softmax")(x)

    # モデルの作成
    model = keras.Model(inputs=inputs, outputs=outputs)
    # モデル構成の表示
    print(model.summary())
    # # モデルの画像保存
    # keras.utils.plot_model(model, "mnist_cnn_classifier.png", show_shapes=True)

    # モデルのパラメータ設定
    model_optimizer = "adam"
    model_loss = "sparse_categorical_crossentropy"
    # 実行パラメータ設定
    num_epochs = 5
    batch_size = 32

    with mlflow.start_run():
        # モデルパラメータを記録
        mlflow.log_param("optimizer", model_optimizer)
        mlflow.log_param("loss", model_loss)
        # 実行パラメータを登録
        mlflow.log_param("num_epochs", num_epochs)
        mlflow.log_param("batch_size", batch_size)

        # ===== オプティマイザ、損失関数、指標を設定してコンパイル
        model.compile(
            optimizer=model_optimizer,
            loss=model_loss,
            metrics=["accuracy"],
        )

        # ===== fitを使ったモデルの訓練
        history = model.fit(
            train_imgs,
            train_labels,
            epochs=num_epochs,
            batch_size=batch_size,
            validation_data=(val_imgs, val_labels),
        )

        # ===== MLflowにトレーニングのメトリクスを記録
        for epoch in range(num_epochs):
            mlflow.log_metric(
                "loss", history.history["loss"][epoch], step=epoch
            )
            mlflow.log_metric(
                "accuracy", history.history["accuracy"][epoch], step=epoch
            )
            mlflow.log_metric(
                "val_loss", history.history["val_loss"][epoch], step=epoch
            )
            mlflow.log_metric(
                "val_accuracy",
                history.history["val_accuracy"][epoch],
                step=epoch,
            )

        # ===== evaluateを使ったテストデータでの評価
        result = model.evaluate(test_imgs, test_labels)
        print(result)

        # ===== predictを使って予測結果を表示
        pred = model.predict(test_imgs)
        print(f"予測: {np.argmax(pred[0])}, 正解: {test_labels[0]}")

        # ===== MLflowにモデルを記録
        # モデルを保存する
        mlflow.keras.log_model(model, "model")


if __name__ == "__main__":
    main()

実装の詳細説明

上記で紹介したプログラムについて MLflow の実装に関わる部分をピックアップして詳細を説明していきます。

インポート

import mlflow
import mlflow.keras

MLflow を利用するためにインポートしている部分です。mlflow.keras は、後述するモデルの保存の際に使用するためにインポートしています。

MLFlow の実験設定

    # ===== MLflowの実験設定
    mlflow.set_experiment("mnist_cnn_classification")

MLFlow を実行する際に今回のプログラムがどういった実験 (Experiment) であるか設定します。設定は mlflow.set_experiment に実験名を文字列で指定します。

このコードは省略することも可能ですが、省略した場合、MLflow は当該プログラムを “Default” 実験として記録します。ただし、複数の実験を行うことが通常だと思いますので、後々で識別をするために実験設定をしておくことを推奨します。

MLflow の実行とパラメータ記録

    # モデルのパラメータ設定
    model_optimizer = "adam"
    model_loss = "sparse_categorical_crossentropy"
    # 実行パラメータ設定
    num_epochs = 5
    batch_size = 32

    with mlflow.start_run():
        # モデルパラメータを記録
        mlflow.log_param("optimizer", model_optimizer)
        mlflow.log_param("loss", model_loss)
        # 実行パラメータを登録
        mlflow.log_param("num_epochs", num_epochs)
        mlflow.log_param("batch_size", batch_size)

上記は、記録するモデルや実行のパラメータを設定して、MLflow を実行している部分です。今回は、オプティマイザーとして "adam"、損失関数はクロスエントロピー "sparse_categorical_crossentropy" を使用しています。また、エポック数は 5、バッチサイズは 32 としています。

MLflow は、mlflow.start_run() を実行して起動し、実験を開始します。start_run を実行すると MLflow は、当該実行が一意になるような Run Name を付与し、後述する MLflow UI で区別することができます。もし、明示的に Run Name を指定したい場合は mlflow.start_run(run_name="My_Run_Name") のように run_name 引数に指定してください。

上記例では start_run を with 句を使用して、その中に実験の内容を記載しています。なお、with 句を使わずに以下のように使うことも可能です。

mlflow.start_run()

# 実験のコードを記載

mlflow.end_run()

with 句を使わない場合は mlflow.end_run() で実験を明示的に終了し、リソースを適切に開放する必要があります。基本的には with 句を使うのが良いでしょう。

MLflow にパラメータを記録するには、mlflow.log_param("optimizer", model_optimizer) のように mlflow.log_param を使用します。第 1 引数にパラメータ名、第 2 引数に記録する値を指定します。

モデルパラメータを探索する場合には、値を変えつつ log_param でしっかり記録をしておきます。これにより、どのパラメータで実験した結果がどういった結果となったが後で確認しやすくなります。

MLflow にトレーニングのメトリクスを記録

        # (省略：モデルのコンパイル、トレーニング実行)

        # ===== MLflowにトレーニングのメトリクスを記録
        for epoch in range(num_epochs):
            mlflow.log_metric(
                "loss", history.history["loss"][epoch], step=epoch
            )
            mlflow.log_metric(
                "accuracy", history.history["accuracy"][epoch], step=epoch
            )
            mlflow.log_metric(
                "val_loss", history.history["val_loss"][epoch], step=epoch
            )
            mlflow.log_metric(
                "val_accuracy",
                history.history["val_accuracy"][epoch],
                step=epoch,
            )

モデルのコンパイルやトレーニングの実行をしたら、結果のメトリクスを記録します。

メトリクスの記録には mlflow.log_metric を使用します。上記例は for 文で各エポックにおけるメトリクスの値を記録するように記載しています。第 1 引数にメトリクス名、第 2 引数に値を指定してます。step 引数にエポックの数値を指定することでどのエポックにおけるメトリクスかを指定しています。

モデルを記録

        # (省略：モデルの評価や予測）

        # ===== MLflowにモデルを記録
        # モデルを保存する
        mlflow.keras.log_model(model, "model")

テストデータでのモデルの評価や予測を行った後に、最後に MLflow にモデルを記録しています。モデルの記録には mlflow.keras.log_model を使用します。これにより実行時のモデルを記録し、バージョン管理や他者との共有などに使うことができます。

なお、モデルの記録のための関数は、scikit-learn のモデルの場合は mlflow.sklearn.log_model、Pytorch の場合は mlflow.pytorch.log_model のように用意されていますので、適切な関数を調べて使用しましょう。

MLflow UI を使った結果確認・管理

プログラム実行後は MLflow UI を使って実行結果の確認等の各種管理ができます。

MLflow UI は上記でも説明した通り「mlflow ui」でサーバを起動したうえで http://127.0.0.1:5000 または http://localhost:5000/ にアクセスします。

MLflow UI にアクセスすると以下のような画面が表示されます。

画面左側の Experiments の部分には、実験が表示されています。プログラムの説明で mlflow.set_experiment にて指定した「mnist_cnn_classification」という実験が表示されます。なお、mlflow.set_experiment を省略した場合は、Default に実験が記録されます。

Experiments にて mnist_cnn_classification を選択したときには、以下のように実験した記録が表示されます。各行がプログラムの 1 回の実行結果だと思ってください。

Run Name には、一意になるように MLFlow が値を設定します。今回上記サンプルコードを実行した 1 回目は「gifted-chimp-416」、2 回目で epoch 数を 100 にして再実行した結果が「indecisive-zebra-710」という Run Name で実行されました。

1 回目の実行結果をクリックした画面を一部抜粋したのが以下になります。

設定したパラメータ情報やメトリックスが表示されていることが分かります。また、登録したモデルは Artifacts の部分に表示されます。

なお、メトリックスの部分のリンクをクリックすると以下のようにグラフを確認することが可能です。

プログラムでは matplotlib などの可視化ライブラリを使ってグラフ化するところを、MLflow UI に任せてしまうことが可能です。

MLflow UI は、他にも実験間の比較をするなど、様々な機能が使えますが今回は概要の紹介にとどめたいと思います。

MLflow におけるリポジトリ (mlruns)

MLflow について紹介してきましたが、MLflow の各種情報が記録されるリポジトリについて簡単に紹介しておきます。

MLflow を使った実行を行うと実行フォルダに「mlruns」というフォルダが生成されることが分かるかと思います。このフォルダが、MLflow のリポジトリとなっていて、各実験データやメタデータ、モデルデータなどを蓄積しています。

mlruns の具体的な構成としては以下のようになります。なお、バージョンによりフォルダ構成が異なる場合があるのでご注意ください。

mlruns/
│
├── 0/                          # 実験ID 0
│   ├── meta.yaml               # 実験0のメタデータ
│   ├── 1234567890abcdef/       # 実行ID
│   │   ├── meta.yaml           # 実行のメタデータ
│   │   ├── params/             # パラメータ
│   │   ├── metrics/            # メトリクス
│   │   ├── tags/               # タグ
│   │   └── artifacts/          # アーティファクト
│   └── ...
├── 1/                          # 実験ID 1
│   └── ...
├── .trash/                     # 削除された実験・実行のデータ
│   └── ...
└── models/                     # モデルレジストリデータ
    ├── MyModelName/            # 登録されたモデル名
    │   ├── 1/                  # モデルのバージョン1
    │   │   ├── artifacts/      # バージョン1のアーティファクト
    │   │   └── meta.yaml       # バージョン1のメタデータ
    │   └── 2/                  # モデルのバージョン2
    │       ├── artifacts/      # バージョン2のアーティファクト
    │       └── meta.yaml       # バージョン2のメタデータ
    └── ...

個々のフォルダの説明は省略しますが、このように MLflow のデータが管理されていることを把握しておいてもらえればと思います。

基本的には実行する Python プログラム配置フォルダに mlruns フォルダがあれば問題ありませんが、異なるフォルダ内にあるリポジトリを使ってプログラムを実行する場合には mlflow.set_tracking_uri() で mlruns フォルダのパスを指定します。

set_tracking_uri を使えば共有フォルダに mlruns を配置しておいて、複数人で共有してリポジトリを使用するといったことも可能です。

まとめ

機械学習ライフサイクル管理のための MLflow の基本的な使い方を解説しました。

MLflow は、機械学習のライフサイクルを管理するためのオープンソースのプラットフォームで、データサイエンティストやエンジニアが機械学習のプロジェクトを効率的に実行できるようになります。

この記事では、MNIST の手書き画像分類を例にして基本的な使い方を紹介してきました。プログラムでパラメータや実験結果、モデルを記録する方法や MLflow UI を使って結果を確認・管理する方法についても紹介しています。

MLflow は、機械学習プロジェクトのライフサイクル管理を容易にしてくれるツールですので、是非使い方を覚えて色々試してもらうと良いかなと思います。

あわせて読みたい

【Python Tech】プログラミングガイド