MLflow モデルをログ記録する

  • [アーティクル]

この記事では、トレーニング済みのモデル (または成果物) を MLflow モデルとしてログする方法について説明します。 MLflow がモデルをパッケージ化する方法とそれらのモデルを実行する方法をカスタマイズするための、さまざまな方法を調べます。

成果物の代わりにモデルをログ記録する理由

「MLflow での成果物からモデルまで」では、成果物またはファイルのログと MLflow モデルのログの違いが説明されています。

MLflow のモデルも成果物です。 ただし、モデルは、モデルを作成した人とそれを使用する人の間のコントラクトとして機能する特定の構造を持っています。 このコントラクトは、成果物自体とその意味の間を橋渡しするために役立ちます。

モデルのログには次の利点があります。

  • mlflow.<flavor>.load_model で推論のためにモデルを直接読み込むことができ、predict 関数を使用できます
  • パイプラインの入力でモデルを直接使用できます
  • スコアリング スクリプトや環境を示さずにモデルをデプロイできます
  • デプロイされたエンドポイントでは Swagger が自動的に有効になり、Azure Machine Learning スタジオでテスト機能を使用できます
  • 責任ある AI ダッシュボードを使用できます

このセクションでは、Azure Machine Learning で MLflow のモデルの概念を使う方法について説明します。

自動ログを使用したモデルのログ記録

MLflow の自動ログ機能を使用できます。 自動ログを使うと、MLflow は使用中のフレームワークに、フレームワークが関連があると見なすすべてのメトリック、パラメーター、成果物、モデルをログするよう指示できます。 自動ログを有効にすると、既定でほとんどのモデルがログされます。 状況によっては、一部のフレーバーでモデルがログされない場合があります。 たとえば、PySpark フレーバーでは、特定のサイズを超えるモデルはログされません。

自動ログをアクティブにするには、mlflow.autolog() または mlflow.<flavor>.autolog() を使います。 次の例では、autolog() を使って、XGBoost でトレーニングされた分類器モデルをログしています。

import mlflow
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score

mlflow.autolog()

model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)

y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)

ヒント

Scikit-Learn パイプラインのような Machine Learning パイプラインを使う場合は、そのパイプライン フレーバーの autolog 機能を使ってモデルをログします。 モデルのログは、パイプライン オブジェクトで fit() メソッドが呼び出されると自動的に行われます。 MLflow を使用した XGBoost 分類器のトレーニングと追跡のノートブックでは、パイプラインを使って前処理を含むモデルをログする方法が示されています。

カスタムのシグネチャ、環境、またはサンプルを使用したモデルのログ記録

MLflow の mlflow.<flavor>.log_model メソッドを使うと、モデルを手動でログできます。 このワークフローでは、モデルのログのさまざまな側面を制御できます。

この方法は次のような場合に使用します。

  • 自動的に検出されるものとは異なる pip パッケージまたはconda 環境を示す。
  • 入力の例を含める
  • 必要なパッケージに特定の成果物を含める
  • autolog ではシグネチャが正しく推論されない。 これは、シグネチャが特定の形状になっている必要があるテンソル入力を処理するときに重要です
  • 何らかの理由で、目的が自動ログの動作でカバーされない

このコード例では、XGBoost 分類器のモデルがログされます。

import mlflow
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score
from mlflow.models import infer_signature
from mlflow.utils.environment import _mlflow_conda_env

mlflow.autolog(log_models=False)

model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
y_pred = model.predict(X_test)

accuracy = accuracy_score(y_test, y_pred)

# Signature
signature = infer_signature(X_test, y_test)

# Conda environment
custom_env =_mlflow_conda_env(
    additional_conda_deps=None,
    additional_pip_deps=["xgboost==1.5.2"],
    additional_conda_channels=None,
)

# Sample
input_example = X_train.sample(n=1)

# Log the model manually
mlflow.xgboost.log_model(model, 
                         artifact_path="classifier", 
                         conda_env=custom_env,
                         signature=signature,
                         input_example=input_example)

Note

  • autolog では、log_models=False という構成が指定されています。 これにより、MLflow モデルの自動ログは行われません。 MLflow モデルの自動ログは、後で手動プロセスとして行われます
  • 入力と出力からシグネチャを直接推論するには、infer_signature メソッドを使います
  • mlflow.utils.environment._mlflow_conda_env メソッドは、MLflow SDK のプライベート メソッドです。 この例では、それによりコードが簡単になりますが、慎重に使ってください。 将来変更される可能性があります。 代わりに、YAML 定義を Python ディクショナリとして手動で生成できます。

predict メソッドで別の動作を使用したモデルのログ記録

mlflow.autolog または mlflow.<flavor>.log_model でモデルをログすると、モデルのフレーバーによって推論の実行方法と、モデルから返される内容が決まります。 MLflow では、predict の結果の生成についてどのような特定の動作も強制されません。 場合によっては、モデルの実行の前と後に、前処理または後処理が必要になることがあります。

このような状況では、入力から出力に直接移動する機械学習パイプラインを実装します。 このような実装は可能であり、パフォーマンス向上のために推奨されることもありますが、実現するのが困難になる可能性があります。 そのような場合は、次のセクションで説明するように、モデルが推論を処理する方法をカスタマイズすると役に立ちます。

カスタム モデルのログ記録

MLflow では、次のような多くの機械学習フレームワークがサポートされています

  • CatBoost
  • FastAI
  • h2o
  • Keras
  • LightGBM
  • MLeap
  • MXNet Gluon
  • ONNX
  • Prophet
  • PyTorch
  • Scikit-Learn
  • spaCy
  • Spark MLLib
  • statsmodels
  • TensorFlow
  • XGBoost

ただし、フレーバーの動作を変更したり、MLflow でネイティブにサポートされないモデルをログしたり、異なるフレームワークの複数の要素を使用するモデルをログしたりすることが必要な場合があります。 このような場合は、カスタム モデルのフレーバーの作成が必要になることがあります。

この問題を解決するため、MLflow には pyfunc というフレーバーが導入されています (Python 関数から開始)。 このフレーバーを使うと、任意のオブジェクトが 2 つの条件を満たしている限り、そのオブジェクトをモデルとしてログできます。

  • 少なくとも predict メソッドを実装している
  • Python オブジェクトが mlflow.pyfunc.PythonModel を継承している

ヒント

Scikit-learn API を実装しているシリアル化可能なモデルでは、モデルが Scikit-learn を使って構築されたかどうかに関係なく、Scikit-learn フレーバーを使ってモデルをログできます。 モデルを Pickle 形式で保存でき、オブジェクトに predict() と predict_proba() メソッドがある場合は、mlflow.sklearn.log_model() を使って MLflow 実行の内部でモデルをログできます。

既存のモデル オブジェクトを囲むラッパーを作成すると、カスタム モデル用のフレーバーの作成が最も簡単になります。 MLflow によって自動的にシリアル化されてパッケージ化されます。 Python オブジェクトは、オブジェクトをファイル システムにファイル (通常は Pickle 形式) として格納できる場合、シリアル化できます。 実行時には、そのファイルからオブジェクトを具体化できます。 これにより、保存時に使用できたすべての値、プロパティ、メソッドが復元されます。

この方法は次のような場合に使用します。

  • Pickle 形式でモデルをシリアル化できる
  • トレーニング直後のモデルの状態を維持したい
  • predict 関数の動作をカスタマイズしたい。

このコード サンプルでは、XGBoost で作成されたモデルをラップして、XGBoost フレーバーの既定の実装とは異なる動作にしています。 クラスの代わりに確率を返します。

from mlflow.pyfunc import PythonModel, PythonModelContext

class ModelWrapper(PythonModel):
    def __init__(self, model):
        self._model = model

    def predict(self, context: PythonModelContext, data):
        # You don't have to keep the semantic meaning of `predict`. You can use here model.recommend(), model.forecast(), etc
        return self._model.predict_proba(data)

    # You can even add extra functions if you need to. Since the model is serialized,
    # all of them will be available when you load your model back.
    def predict_batch(self, data):
        pass

実行でカスタム モデルをログします。

import mlflow
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score
from mlflow.models import infer_signature

mlflow.xgboost.autolog(log_models=False)

model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
y_probs = model.predict_proba(X_test)

accuracy = accuracy_score(y_test, y_probs.argmax(axis=1))
mlflow.log_metric("accuracy", accuracy)

signature = infer_signature(X_test, y_probs)
mlflow.pyfunc.log_model("classifier", 
                        python_model=ModelWrapper(model),
                        signature=signature)

ヒント

この infer_signature メソッドは、y_probs を使ってシグネチャを推論します。 ターゲット列にはターゲット クラスがありますが、モデルでは各クラスの 2 つの確率が返されるようになりました。

次のステップ