MLflowの使い方 - 機械学習初心者にもできる実験記録の管理 -

こんにちは！nakamura（@naka957）です。今回はMLflowをご紹介します。

読者の皆さんは、機械学習モデルのハイパーパラメータ調整を手作業で管理し、苦労した経験がないでしょうか。実験記録の管理は大事な一方で、なかなか大変です。

今回紹介するMLflowは、実験記録を簡単に管理できる便利なPythonライブラリです。MLflowは実験管理だけでなく、機械学習プロジェクト全体を管理する様々な機能を提供する非常に人気なライブラリです。一方で、多機能な反面で初心者が最初に導入するにはハードルが高い側面があるのも事実です。

本記事では、MLflowの実験管理の機能に絞り、簡単な例で使い方をご説明します。そのため、初めて使用する方も安心してご覧ください。

では、早速始めていきます。

実験記録の重要性
MLflowとは
MLflowのインストール
データセット準備
機械学習モデルの用意
MLflowを用いた実験記録の管理
実験記録をUIで可視化
まとめ
参考文献

実験記録の重要性

実験記録を残すことは重要です。実験記録を残さないと、過去の実験から最良の選択ができないことに加え、再現することもできません。特に、ビジネス上で機械学習を活用する場合、運用モデルで生じた問題をすばやく再現し、検証することは必須となります。

一方で、機械学習は様々な要素で構成されるため、実験記録の管理が煩雑になります。例えば、使用するデータセット、前処理方法、機械学習アルゴリズム、ハイパーパラメータ条件など、多くの要素が挙げられます。

私自身もモデルのハイパーパラメータ調整で実験記録を取る機会が頻繁にありますが、下図のような表形式で管理し、苦労した苦い経験が多々あります。

MLflowとは

MLflowはコードの実行ごとに実験記録を作成してくれる便利なライブラリです。OSSのため、誰でも無料で使用することができます。

MLflowは実験管理だけでなく、機械学習プロジェクト全体で使用できる機能を提供しており、4つの主要な機能（コンポーネント）で構成されています。

MLflow Tracking: 実験記録の管理
MLflow Projects: 再利用可能な形でコードをパッケージ化
MLflow Models: 本番環境へ機械学習モデルを提供するための機能
Model Registry: 運用管理のための、モデル登録・API・UIの提供

今回は、機械学習モデル構築時の実験記録の管理で使用する、MLflow Trackingの使い方を実際の例を通してご紹介します。具体的には、機械学習モデルのハイパーパラメータの水準とその評価指標の実験記録をMLflowで管理します。また、実験結果をUIで可視化して確認します。

MLflowの導入で実験条件とその結果（評価指標、学習済みモデル）を効率よく管理できるようになり、負担の軽減に繋がることを実感頂けるはずです。

では、実際に使っていきましょう。

MLflowのインストール

以降で紹介するコード例は、Jupyter Notebookでの実行を想定しています。また、本記事のサンプルコードは、次のライブラリのバージョンで動作確認をしています。

scikit-learn==1.0.2
mlflow==1.25.1

MLflowのインストールが未実施の人は、事前にインストールが必要です。

scikit-learn==1.0.2
mlflow==1.25.1

インストールの完了後、使用するライブラリを事前に読み込んでおきましょう。

from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error as MSE
from sklearn.metrics import r2_score

import mlflow
import mlflow.sklearn

データセット準備

今回はscikit-learnライブラリに含まれるカリフォルニア住宅価格のデータセットを使用します。

データセットはsklearn.datasetsモジュール内のfetch_california_housing()から取得できます。引数に”as_frame=True”を指定すれば、dataset['frame']からpandasのDataFrame形式で取得できます。

# as_frame=True --> 'frame'にpandasのDataFrame形式で格納される
dataset = fetch_california_housing(as_frame=True)
df = dataset['frame']

df.head()

データセットの１行分に相当する各データは、地域をブロックごとに分けて算出された値で構成されています。例えば、目的変数の変数MedHouseValは対象ブロックの住宅価格の中央値を表します。同様に、説明変数も各ブロックごとで算出された値です。

各変数の詳細は次の通りです。

MedInc: 世帯ごとの所得の中央値
HouseAge: 築年数の中央値
AveRooms: 1世帯当たりの部屋数の平均値
AveBedrms: 1世帯当たりの寝室数の平均値
Population: 人口
AveOccup: 世帯人数の平均値
Latitude: 対象ブロックの中心緯度
Longitude: 対象ブロックの中心経度
MedHouseVal: 住宅価格の中央値

データセットを説明変数と目的変数に分けます。目的変数は変数MedHouseValです。

target_col = 'MedHouseVal'

X, y = df.drop(columns=[target_col]), df[target_col]

次に、データセットを訓練用と検証用に分割します。

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)

以降で、訓練データで機械学習モデルを学習し、評価指標を検証データから算出し、モデルの精度を確認していきます。

機械学習モデルの用意

今回はランダムフォレストを機械学習モデルに使用します。scikit-learnライブラリで簡単に使用することができます。

次のコードで、モデルの定義と訓練を行います。今回、モデルの引数に指定するハイパーパラメータは次の通りです。後ほど、最初の2つ（n_estimators, max_depth）をMLflowで管理していきます。

n_estimators: 決定木の数
max_depth: 決定木の深さ
criterion: 評価指標
random_state: 再現性確保のための乱数シード値

# モデルの定義
model = RandomForestRegressor(
    n_estimators = 10,
    max_depth = 5,
    criterion = 'squared_error',  # 'mse'
    random_state = 0,
)

# 訓練
model.fit(X_train, y_train)

訓練したモデルで検証データの精度を確認してみましょう。評価指標は平均二乗誤差であるmean squared error(MSE)とします。

# 検証データの予測
y_pred = model.predict(X_test)

# MSE誤差の算出
mse = MSE(y_pred, y_test)
print(f"MSE = {mse:.3}")  # 小数点以下３桁まで出力


>>  MSE = 0.487

以上で、機械学習モデルを用意し、その評価指標を確認しました。

本記事の目的は、モデルのハイパーパラメータの変更とその評価指標の実験管理です。その準備として、複数水準のハイパーパラメータで実行するコードへ拡張しましょう。

まず、試行1回に相当する先ほどのコード（モデル定義、訓練、検証データの予測、評価指標の算出）を1つの関数で定義します。ただし、ハイパーパラメータのn_estimatorsとmax_depthは試行ごとに変更できるように関数の引数とします。また、実験結果となる評価指標を関数の返り値にします。練習のために、評価指標にはMSEだけでなく、R2値（決定係数）も含めてみましょう。R2値は予測値と実データの一致の程度を表す指標で、値が1に近いほどモデル精度が良いと期待できます。

# 試行関数
def train(n_estimators, max_depth):
    # モデルの定義
    model = RandomForestRegressor(
        n_estimators = n_estimators,
        max_depth = max_depth,
        criterion = 'squared_error',  # 'mse'
        random_state = 0,
    )

    # 訓練
    model.fit(X_train, y_train)

    # 検証データの予測
    y_pred = model.predict(X_test)

    # MSE誤差, R2値の算出
    mse = MSE(y_pred, y_test)
    R2 = r2_score(y_pred, y_test)

    return model, mse, R2

次に、ハイパーパラメータを複数水準用意し、試行関数を用いて各水準での精度を順に確認してみましょう。

# ハイパーパラメータの水準
cand_n_estimators = [10, 100, 1000]
cand_max_depth = [1, 5, 10]

trial = 0
for n_estimators in cand_n_estimators:
    for max_depth in cand_max_depth:
        trial += 1
        model, mse, R2 = train(n_estimators, max_depth)
        # 小数点以下３桁まで出力
        print(f"trial {trial}: n_estimators={n_estimators}, max_depth={max_depth}, MSE = {mse:.3}, R2 = {R2:.3}")


>>  trial 1: n_estimators=10, max_depth=1, MSE = 0.91, R2 = -1.33
>>  trial 2: n_estimators=10, max_depth=5, MSE = 0.487, R2 = 0.38
>>  trial 3: n_estimators=10, max_depth=10, MSE = 0.312, R2 = 0.676
>>  trial 4: n_estimators=100, max_depth=1, MSE = 0.909, R2 = -1.32
>>  trial 5: n_estimators=100, max_depth=5, MSE = 0.481, R2 = 0.382
>>  trial 6: n_estimators=100, max_depth=10, MSE = 0.304, R2 = 0.679
>>  trial 7: n_estimators=1000, max_depth=1, MSE = 0.905, R2 = -1.33
>>  trial 8: n_estimators=1000, max_depth=5, MSE = 0.479, R2 = 0.382
>>  trial 9: n_estimators=1000, max_depth=10, MSE = 0.302, R2 = 0.681

ハイパーパラメータ（n_estimators, max_depth）の各水準でのモデル精度が確認できました。

以上で、MLflowを用いて実験管理をする準備は整いました。上記ではprintの出力で結果を確認しましたが、MLflowではUIから結果を確認することができます。

MLflowを用いた実験記録の管理

前節のコードにMLflowを用いた追加変更を行います。数行のコードを書き換えるだけで適用できます。

以下の変更後のコードで追加変更した箇所は、コード内の「""" 追加箇所 """」になります。

大事な点は、MLflowの使用を明示する”with mlflow.start_run():”と、記録を取る変数（実験条件と評価指標）をMLflowの機能で指定することです。

""" 追加箇所 """
# コード実行に対する管理IDを発行
with mlflow.start_run():
    # ハイパーパラメータの試行水準
    cand_n_estimators = [10, 100, 1000]
    cand_max_depth = [1, 5, 10]

    trial = 0
    for n_estimators in cand_n_estimators:
        for max_depth in cand_max_depth:
            """ 追加箇所 """
            # 各水準それぞれの管理IDを発行
            with mlflow.start_run(nested=True):
                trial += 1
                model, mse, R2 = train(X_train, y_train, X_test, y_test, n_estimators, max_depth)
                # 小数点以下３桁まで出力
                print(f"trial {trial}: n_estimators={n_estimators}, max_depth={max_depth}, MSE = {mse:.3}, R2 = {R2:.3}")

                """ 追加箇所 """
                # ハイパーパラメータ, 評価指標, 学習済みモデルをMLflowへ保存 
                mlflow.log_param("n_estimators", n_estimators)
                mlflow.log_param("max_depth", max_depth)
                mlflow.log_metric("mse", mse)
                mlflow.log_metric("R2", R2)

                mlflow.sklearn.log_model(model, "model")

なお、上記コードで、”with mlflow.start_run():”が2箇所ありますが、これらは、1箇所目でコード実行に対応する管理IDを発行し、2箇所目でハイパーパラメータの各水準での管理IDを発行する形で実験記録を保存しています。

問題なく実行が完了すれば、実行結果が保存されます。実験記録の保存場所は、実行フォルダ内に作成される、”./mlruns”フォルダです。ただし、直接フォルダ内のファイルを確認するわけではなく、次節で紹介するように、UI上で簡単に確認できるところもMLflowの便利な点です。

実験記録をUIで可視化

前節の実験記録をUIで可視化し、確認してみましょう。MLflowにはUIが含まれており、コマンド1つで立ち上げることができます。

次のコマンドを、Windowsの場合はコマンドプロンプト上で、Macの場合はターミナル上で実行してください。

$ mlflow ui

# port番号を指定する場合  ex. port番号5001の場合
$ mlflow ui --port 5001

MLflowのUIはWebブラウザ上からアクセスできます。port番号を指定しない場合、デフォルトで5000が割り当てられるため、”http://localhost:5000”でアクセスできます。また、port番号を指定した場合は、指定したport番号でURLにアクセスしてください。

問題なく起動すれば、次のようなUIが立ち上がります。

早速、先ほどの実行結果を確認してみましょう。実験記録が保存されていることが確認できます。

下図に、上図右側を拡大した図を示します。評価指標（R2, MSE）とハイパーパラメータ（n_estimators, max_depth）の水準が表示されています。また、例えば、R2箇所をクリックすることで、R2値に基づき結果をソートすることも可能です。このように、評価指標に基づく実験結果の確認とそのハイパーパラメータを追跡でき、実験結果の再現が容易になります。