Databricksを活用したMLOps入門：MLFlowのSignature機能でモデル管理を簡単に

2024年12月23日 09:05

三菱UFJフィナンシャル・グループ（以下MUFG）の戦略子会社であるJapan Digital Design（以下JDD）でMUFG AI Studio（以下M-AIS）に所属する小林です。普段はMUFG各社向けにAI施策のPoCやモデル実装等を担当しています。

今回は普段モデル運用を行う中で利用する機会が多いMLFlowのSignatureについて紹介いたします。弊社では普段からDatabricksを活用し、モデル構築および実装、本番運用を手掛けております。

MLFlowについての説明はこちらに譲る事として、本記事では前半でMLFlowにおけるSignatureについて説明し、後半で実装方法および注意点を述べます。

Signatureとは？

MLFlowのSignatureは、モデルの入出力スキーマ（データ型や構造）を明示的に定義するための便利な機能です。この機能を活用することで、MLエンジニアをはじめとするモデル実装者がデータサイエンティストから引き継いだコードを理解する際に抱く、以下のような疑問に対する答えを明確に示すことができます：

どのような形式の入力が必要か？
どのような形式の出力が期待されるか？

Signatureを活用することで、モデルの使用方法が明確になり、運用時のトラブルを未然に防ぐことが可能になります。

なぜSignatureが重要なのか？

機械学習モデルを運用環境で活用する際、以下のような課題が頻繁に発生します：

データ型の不整合
モデルに期待されるデータ型（例：整数や文字列）が実際に与えられたデータと一致しない場合、予期しないエラーが発生することがあります。
構造の誤解
入力データの列名や順序が正しく提供されない場合、一般的にはモデルの出力結果は信頼できません。
再現性の低下
モデルのトレーニング後に正しい入力データの形式に関する知識や情報が残されずに失われてしまうと、再トレーニングやデバッグが困難になります。

Signatureのメリット

Signatureを使用することで、これらの課題を効果的に解決できます。具体的には、以下のような利点があります：

モデルの信頼性向上：期待するデータ形式が明確になるため、エラーを減らせます。
再現性の確保：トレーニング時と同じ条件でモデルを再利用できるようになります。
保守性の向上：チーム内での共有や運用時のトラブルシューティングが容易になります。

Signatureのコンポーネント

MLFlowのSignatureは、以下の3つのスキーマタイプで構成されています。それぞれの役割を理解することで、より効果的にSignatureを活用できます。

1. Input（入力スキーマ）

役割：モデルが受け取る入力データの型や構造を定義します。
用途：データ型のチェックを通じて、予期しないエラーの発生を防ぎます。
例）列名、型、順序が正しいか確認するための基準。

2. Output（出力スキーマ）

役割：モデルが生成する出力データの型や構造を定義します。
用途：モデルが期待どおりの出力を返しているか検証します。

3. Params（パラメータスキーマ）

役割：推論時にオプションで指定可能なパラメータを定義します。
用途：特定の条件や柔軟な処理を実現するためのオプションを提供します。

ポイント：
Paramsをうまく活用することで、予測処理を柔軟にカスタマイズすることが可能です。たとえば、特定の閾値を指定して出力を調整するなど、さまざまな運用シナリオに対応できます。

Signatureの付与方法

MLFlowでSignatureを付与する際の具体例を示します。この例では、LightGBMを使用したモデルのトレーニングとSignatureの設定を行います。

前提環境

DatabricksRuntime：15.4 LTS
mlflow：2.15.1
lightGBM：4.3.0

ModelWrapperクラスの定義

MLFlowでは、Signatureを付与してモデルをロギングする場合、必ずしもpyfunc flavourでロギング(mlflow.pyfunc.log_model)する必要はありません。しかし、保守性を高めるためには、モデルをpyfunc flavourでロギングすることを個人的にはおすすめします。

基本的にpyfunc flavourでモデルを保存した場合、予測処理はpredict関数しか実行出来ませんが、以下例のように、predict_probaやpredict_log_probaといった複数の予測メソッドを柔軟に実装する事が可能となります。

さらに、pyfunc flavourを利用することで、予測時にSHAP値を計算し、予測結果に影響を与えた特徴量のSHAP値を含めるといった柔軟な処理も可能です。こうした柔軟性は、pyfunc flavourでモデルをロギングする大きなメリットだと考えます。

import mlflow.pyfunc
from joblib import load
class ModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self,model):
        self.model = model
    def predict(self, context, model_input, params=None):
        # デフォルトの予測方法を設定
        params = params or {"predict_method": "predict"}
        predict_method = params.get("predict_method")
        # 指定された予測方法に応じてモデルの予測を実行
        if predict_method == "predict":
            return self.model.predict(model_input)
        elif predict_method == "predict_proba":
            return self.model.predict_proba(model_input)
        elif predict_method == "predict_log_proba":
            return self.model.predict_log_proba(model_input)
        else:
            raise ValueError(f"The prediction method '{predict_method}' is not supported.")

以下に実際の学習時のコードを記載いたします。

import mlflow
from joblib import dump
import lightgbm as lgb
from mlflow.models.signature import infer_signature
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
import pandas as pd
# Breast Cancerデータセットをロード
data = load_breast_cancer()
X = pd.DataFrame(data.data, columns=data.feature_names)
y = pd.Series(data.target, name="target")
# 学習用とテスト用にデータセットを分割
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)
# LightGBM用にデータセットを作成
dtrain = lgb.Dataset(X_train, label=y_train)
dtest = lgb.Dataset(X_test, label=y_test, reference=dtrain)
# 学習パラメータを設定
params = {"objective": "binary", "metric": "binary_logloss", "verbosity": -1}
# コールバック関数を使用して早期停止を設定
callbacks = [lgb.early_stopping(stopping_rounds=10)]
# モデルの学習
with mlflow.start_run() as run:
    lgbm_model = lgb.train(
        params,
        dtrain,
        num_boost_round=100,
        valid_sets=[dtrain, dtest],
        valid_names=["train", "valid"],
        callbacks=callbacks,
    )
    # シグネチャを計算
    signature = infer_signature(
        X_train,
        lgbm_model.predict(X_train),
        params={"predict_method": "predict"},
    )
    # モデルをMLflowにログ
    mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=ModelWrapper(lgbm_model),
        signature=signature,
    )
    run_id = run.info.run_id
    model_uri = f"runs:/{run_id}/model"
    print(f"Model saved in run: {run_id}")

ポイントとしては以下のsignatureを設定する場合の処理およびMLFlowに登録する際の箇所です。

以下の箇所ではinput,output,paramsを取得するためのコードとなります。

（上述したコードの再掲）
    # シグネチャを計算
    signature = infer_signature(
        X_train,
        lgbm_model.predict(X_train),
        params={"predict_method": "predict"},
    )

また、以下の箇所では実際に作成したモデルに対してsignatureを付与し保存する箇所となります。

（上述したコードの再掲）
    # モデルをMLflowにログ
    mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=ModelWrapper(lgbm_model),
        signature=signature,
    )

予測時におけるSignatureの挙動

続いて実際のSignatureの挙動を紹介いたします。

まずsignatureではカラム名まで覚えているため、カラムの順序がシャッフルされても自動で元の順序に直し、予測処理を出してくれます。

loaded_model = mlflow.pyfunc.load_model(model_uri)
# 特徴量の順序を変更したテストデータを作成
X_test_shuffled = X_test[X_test.columns[::-1]]
# 学習時と同じ順序で予測
predictions_proba_1 = loaded_model.predict(X_test)
print("Non Shuffle Prediction Probabilities:", predictions_proba_1[:10])
# 順序を変更した特徴量で予測確率を取得
predictions_proba_2 = loaded_model.predict(X_test_shuffled)
print("Shuffle Prediction Probabilities:", predictions_proba_2[:10])
# 特徴量シャッフル有無で予測結果が同一か
assert predictions_proba_1.sum() == predictions_proba_2.sum()

また、signatureの真の効果としてカラム型が異なる場合エラーとなります。

# 学習する型を変更した処理
X_test_error = X_test
X_test_error["mean radius"] = X_test_error["mean radius"].astype("int")
print(X_test_error.dtypes)
predictions_proba = loaded_model.predict(X_test_error)

Signature実装時の注意事項

Signatureを実際に活用する際には、いくつか注意すべきポイントがあります。今回はLightGBMを利用したケースを踏まえた注意点を提示いたします。

最後にlightGBMを利用する場合APIが複数種類存在し、log_modelやload_modelにもAPIが以下のように複数種類存在するため、それぞれの挙動をまとめました。

lightGBMのAPI
- TrainingAPI
- SklearnAPI
log_modelのAPI
- mlflow.lightgbm.log_model
- mlflow.sklearn.log_model
- mlflow.pyfunc.log_model
load_modelのAPI
- mlflow.lightgbm.load_model
- mlflow.sklearn.load_model
- mlflow.pyfunc.load_model

注目すべきポイント

MLFlowでSignatureを活用する際に、特に注目していただきたい重要なポイントは以下の2点です：

1. Signatureを有効化するためのモデルのロード方法

Signatureを有効化するためには、モデルのロード時にmlflow.pyfunc.load_modelを使用する必要があります。
たとえば、mlflow.lightgbm.load_modelやmlflow.sklearn.load_modelを使用した場合、Signatureが有効化されないため注意が必要です。

2. Signatureを使用しない場合の予測時のリスク

mlflow.sklearnやmlflow.lightgbmを使用してlog_modelとload_modelを行う場合の注意点：

カラムの順序が変わっても予測が実行されてしまう
カラムのデータ型が変わっても予測が実行されてしまう

これにより、意図しないデータで予測が行われ、誤った結果を生む可能性があります。この挙動に気づかない場合、バグの原因となり、運用上の大きな問題を引き起こす恐れがあります。

特に2点目については気づかずバグが発生してしまう温床となるため、正しい使い方を理解したうえで利用しましょう。

以上のポイントを踏まえ、Signatureを利用することで以下のような利点があります：

データの順序や型を保証することで、予測の信頼性を担保できる。
誤ったデータ形式の検知が可能になり、潜在的なバグを防止できる。

MLFlowのSignatureは、簡単に導入できるだけでなく、予測時の挙動を頑健にするため、MLOpsにおける非常に有用な機能です。
もしご興味がありましたら、ぜひ一度お試しいただき、運用環境での適用をご検討ください！

Japan Digital Design株式会社では、一緒に働いてくださる仲間を募集中です。カジュアル面談も実施しておりますので下記リンク先からお気軽にお問合せください。

この記事に関するお問い合わせはこちら

Japan Digital Design 株式会社
MUFG AI Studio
Yuto Kobayashi