MLflow

【MLflow】基本的な使い方 ~機械学習のライフサイクル管理~

【MLflow】基本的な使い方 ~機械学習のライフサイクル管理~

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

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

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

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

機能概要
MLflow Tracking実験のパラメータ、コード、メトリクス、結果を記録する機能です。ユーザーは異なる実験結果を比較し、最適なモデルを選択できます。
MLflow Projectsプロジェクトの構造や依存関係を定義するための形式です。これによりコード、データ、環境設定を一元管理して、再現性と共有の容易さを向上させます。
MLflow Models異なる機械学習フレームワークからのモデルを一つの標準フォーマットで保存、再利用、共有するための機能です。これにより、異なる環境へのモデルのデプロイが容易になります。
MLflow Model Registryモデルのバージョン管理、ライフサイクルのステージング、注釈の付与などができる中央リポジトリ。モデルの使用状況を追跡して運用環境での利用を管理する

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

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

MLflowの環境準備

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

MLflowのインストール

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

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の実装に関わる部分をピックアップして詳細を説明していきます。

CNNに関わる部分の説明は省略します。「CNN(畳み込みニューラルネットワーク)による画像分類の基本」を参考にしてください。

インポート
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を起動する必要がありますが、その際には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を使用します。第一引数にパラメータ名、第二引数に記録する値を指定します。

モデルパラメータを探索する場合には、値を変えつつ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文で各エポックにおけるメトリクスの値を記録するように記載しています。第一引数にメトリクス名称、第二引数に値を指定してます。step引数にエポックの数値を指定することでどのエポックにおけるメトリクスかを指定しています。

モデルを記録
        # (省略:モデルの評価や予測)

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

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

なお、モデルの記録のための関数は例えばScikit-Learnのモデルの場合は、mlflow.sklearn.log_model、Pythorchの場合は、mlflow.pytorch.log_modelといったように用意されていますので、適切な関数を調べて使用するようにしましょう。

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

紹介したプログラムを実行した後には、 MLflow UIを使って結果を確認したり、管理したりすることができます。

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

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

Note

MLflowの公式ドキュメントはこちらを参照してください。