はじめに

ホクソエムサポーターの藤岡です。会社を移りましたが、相変わらずPythonを書く仕事をしています。

前回の記事に引き続き、今回もMLflowについての記事です。前回はトラッキング寄りでしたが、今回はモデルのデプロイにも関わってくる内容です。

MLflowはXGBoost, PySpark, scikit-learnといった多様なライブラリに対応していて、様々な機械学習タスクに活用することができるのが売りの一つです。その実現のため、設計や実装に様々な工夫がされているのですが、この部分について詳しくなることで、オリジナルの機械学習モデルをMLflowとうまく繋ぐことができるようになったり ETLのようなモデル学習にとどまらない使い方もできるようになったりします。

本記事では、XGBoostをMLflowで扱うためのモジュール mlflow.xgboost について解説することで、拡張性の鍵であるflavorについて紹介します。また、その前提知識として、MLflow Model, PyFunc Modelについても関連する部分をさらっと解説します。

機械学習が完全に絡まない部分で使う必要があるのかはさておき、機械学習が絡んでいるフローを広く一つのフレームワークで管理できるようにしておくのは有用だと思います。実際、MLflowのサンプルコードにもETLを含めて管理している例があります*1。

本記事を通してMLflowについて理解を深め、活用の幅を広げましょう！

なお、本記事はMLflowの使用方法と前回記事の内容 (のうちMLmodel周り) が前提知識となります。また、MLflow v1.10.0 時点での内容です。

MLflow Model

今日では、日々新しいMLモデルが提案され、それらを実装したパッケージも次々と生まれています。それらはsklearnの提供する形式に従っているものから、 xgboostのようなデータ構造まで自前で用意してパッケージングしたものまで様々であり、統一的な形式のようなものはありません。

MLflowは抽象的なモデルフォーマットを定義し、それを通じてこれらのフレームワークやパッケージを統一的なインターフェースで扱います。このモデルフォーマットは MLflow Model と呼ばれます。モデルとは言っても機械学習モデルというよりそれにまつわる種々の内容を定めたものです。以下の例をはじめとする種々の項目についてそれぞれ形式を定めています。

ストレージ
モデルの入出力形式（Signature）
モデルのAPI （セーブ、ロード、ロギング等）

このフォーマットはmlflow.models以下で実装されています。例えば、本体となるクラスはmlflow.models.model.pyの中にModel()クラスとして定義されています。ちなみに、このクラスがどんなメソッドやパラメータを持っているか一通り目を通しておくと、MLflow Modelについて少しイメージがしやすくなります。

flavor

MLflow Modelは抽象度が高く、それ単体では使うことができません。実際、モデルのデプロイ部分の実装を見てみると、入力とメソッド名のみを定義している以外はすっからかんとなっています。

そのため、MLflow Modelは既存の機械学習フレームワーク等のモジュールやスクリプトがこのフォーマットに従うようにするためのインターフェースを必要とします。これを、flavorと呼びます。例えば、xgboost.Booster (xgboostのモデルクラス) をMLflowで扱う場合、Boosterに特定のメソッドを生やすためのラッパなどをflavorとして実装すれば、 MLflowの学習トラッキングやモデルデプロイなどの機能をxgboost.Boosterに対して実行できます。

とはいっても、xgboost.Boosterについてはmlflow.xgboostモジュールとしてflavorが既に定義されているので、ユーザ側ではただmlflow.xgboostの関数を呼び出すだけで使えます。

さらに、あるflavorをベースに別のflavorを実装することも可能です。ベースにするといっても、実際には両方を併用するような形になります。

flavorの種類

flavorは大きく3種類に分けることができます。

Built-in flavor

XGBoostを始めとして、MLflowが対応している種々のモジュールに対してそれぞれflavorが定義されています。一般的なデータ分析タスクで使うような機械学習モデルはたいていの場合はこれらのflavorで対応ができるので、小さいコーディングコスト（log_model() と save_mode()を呼び出すだけ）でMLflowの提供する機能を十全に使うことができます。

さらに、いくつかのflavorにはautolog機能を備えたflavorがあり、その場合はほぼ全自動でMLflowの機能を使うことができます。

PyFunc Model flavor

PyFunc Modelとは、特定の機械学習モデルを指すモデルではなく、以下の入出力をもつ関数の抽象モデルです。

predict(model_input: pandas.DataFrame) -> [numpy.ndarray | pandas.(Series | DataFrame)]

これ自身を使うのではなく、これをベースに別のflavorを作るために用意されています。上記の入出力をもつ関数として表現可能なものであればこのflavorを元に作成可能です。全てのBuilt-in FlavorはPyFunc Model Flavorをベースに作成されています。

MLflowのdocumentやソースコードを読んでいるとき、PyFunc ModelとMLflow Modelを混同しやすいので注意しましょう。基本的にはPyFunc ModelがMLflow Modelの一部だとみて問題ないとは思いますが、とにかくmodelという単語がたくさん出てきて混乱するので区別を付けられるようになっておくのがいいです。

ここでは紹介しませんが、PyFunc Modelの詳細について知りたい方はこちらやこちらが詳しいです。 backendの実装等を覗いてみると、PyFunc Model flavorが非常に多くの機能を実装し、その他のflavorの作成を助けているのかが見て取れておもしろいです。

Custom Python flavor

MLflowが提供しているflavor以外にもユーザ側がflavorを開発することでMLflow互換のモデルを定義することが可能です。

作成方法は以下の2通りがあります。

実装したいモデルのクラスをmlflow.pyfunc.PythonModelのサブクラスとして作成
既存のモデルをPyFunc Modelと相互変換するような入出力関数を定義

基本的には1が簡単かつ有用ですが、MLflowを使わないで保存したモデルオブジェクトとの互換性はないので、既に学習済みのモデルを取り込むなどの互換性が必要な場合は2の方法を採る必要があります。

XGBoost flavor

Overview

f:id:kazuya_fujioka:20200817090002p:plain — XGBoost Flavorとその関連する要素

図は、XGBoost flavor, MLflow, XGBoostの三者、そしてflavorを呼び出すスクリプトの4要素がどのオブジェクト (関数/クラス/モジュール) を通じて繋がっているのかを図示したものです。黄色い四角で囲われた6つの関数がXGBoost flavorの主な構成要素です((autolog()等は本質ではなく、理解も容易なので省略しています。))。なお、分かりやすさのためPyFunc Modelの実装であるpyfuncモジュールもMLflowに入れています。

ユーザ側はlog_model()（MLflow Modelの記録）、load_model()（モデルの読み込み）を操作します。内部的には、log_model()の呼び出しでsave_model()（MLflow Modelの書き出し）が内部的に呼び出されるほか、_load_pyfunc()はモデルのデプロイ時に内部的に呼び出されます。

これらは、基本的には全てのflavorに共通です。

一方で_load_model()と_XGBModelWrapper()はxgboostライブラリに固有の実装です。 XGBoost flavor内から必要に応じて呼び出されます。

Components

では、これらを順を追って詳細に読んでいきます。なお、解説のために動作が変化しない程度にコメントや変数等を変更している部分があります。

_load_pyfunc() & _load_model() & _XGBModelWrapper

def _load_pyfunc(path):
    xgb_original = _load_model(path)
    return _XGBModelWrapper(xgb_original)


def _load_model(path):
    import xgboost as xgb
    model = xgb.Booster()
    model.load_model(os.path.abspath(path))
    return model


class _XGBModelWrapper:
    def __init__(self, xgb_model):
        self.xgb_model = xgb_model

    def predict(self, dataframe):
        import xgboost as xgb
        return self.xgb_model.predict(xgb.DMatrix(dataframe))

これらの3つの関数で実現していることは、引数で与えられたパスpathにある学習済みxgboostモデルをロードする処理を、 _load_pyfunc()という名前の関数として実装することです。

_load_pyfunc()関数はpyfunc.__init__.py内に定義されているload_model()関数内で呼び出され、 PyFuncModelクラスの初期化に使用されます (実装)。

モデルのロード部分はxgboostライブラリの機能をそのまま使って_load_model()関数として実装しています。ただし、xgboost.Booster()はpredict()メソッドを実装しているものの、その入力がPandas DataframeではなくDMatrixなのでそのままではPyFuncModelに渡せません。その間を埋めるためのラッパとして_XGBModelWrapperクラスが定義されています。

save_model()

def save_model(xgb_model, path, conda_env=None, mlflow_model=None,
               signature: ModelSignature=None, input_example: ModelInputExample=None):
    import xgboost as xgb

    path = os.path.abspath(path)
    if os.path.exists(path):
        raise MlflowException("Path '{}' already exists".format(path))
    os.makedirs(path)
    if mlflow_model is None:
        mlflow_model = Model()
    if signature is not None:
        mlflow_model.signature = signature
    if input_example is not None:
        _save_example(mlflow_model, input_example, path)
    model_data_subpath = "model.xgb"
    model_data_path = os.path.join(path, model_data_subpath)

    # Save an XGBoost model
    xgb_model.save_model(model_data_path)

    conda_env_subpath = "conda.yaml"
    if conda_env is None:
        conda_env = get_default_conda_env()
    elif not isinstance(conda_env, dict):
        with open(conda_env, "r") as f:
            conda_env = yaml.safe_load(f)
    with open(os.path.join(path, conda_env_subpath), "w") as f:
        yaml.safe_dump(conda_env, stream=f, default_flow_style=False)

    pyfunc.add_to_model(mlflow_model, loader_module="mlflow.xgboost",
                        data=model_data_subpath, env=conda_env_subpath)
    mlflow_model.add_flavor(FLAVOR_NAME, xgb_version=xgb.__version__, data=model_data_subpath)
    mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))

MLflow Modelmlflow_modelをartifact化し、渡されたローカルのフォルダパスpathに保存する関数です。具体的には、フォルダの中に以下のファイルが作成されます。

signatureとして渡されたSignatureのダンプファイル
xgb_modelとして渡されたxgboost.Boosterのダンプ。
conda_envとして渡された辞書/ファイル名のconda envファイル（ただし、無ければデフォルトを生成）
MLflow Modelファイル（MLmodel）

最も重要なのは、以下の3行です。

pyfunc.add_to_model(mlflow_model, loader_module="mlflow.xgboost",
                        data=model_data_subpath, env=conda_env_subpath)
mlflow_model.add_flavor(FLAVOR_NAME, xgb_version=xgb.__version__, data=model_data_subpath)

MLmodelファイル（yaml）にはflavorごとの設定項目を記録するセクションがあり、その内容をここで生成しています。

あるflavorの設定項目を記録するには、Model.add_flavor()メソッドを使って、 flavor名（第一引数）とその設定項目（キーワード引数）を渡します。

例えば、3行目で、XGBoost flavorの項目が以下のように作成されます。

xgboost:
    xgb_version: <xgb.__version__の値>
    data: subpath/from/MLmodel/root/to/model.xgb

なお、セクション名"xgboost"はmlflow.xgboost内で

FLAVOR_NAME = "xgboost"

と定義されているところから来ています。

さらに、PyFunc Model flavorをベースに作られているflavorを使った場合、 PyFunc Modelのセクション"python_function"を作成します。

1-2行目の"pyfunc.add_to_model()"関数は"Model.add_flavor()"メソッドの呼び出しをPyFunc Model flavor向けにラップしたものです。この中で conda env等のPyFuncの基本的な構成要素へのパスとPythonのバージョンを含んだ"python_function"セクションがmlflow_modelに記録されます。

log_model()

def log_model(xgb_model, artifact_path, conda_env=None, registered_model_name=None,
              signature: ModelSignature=None, input_example: ModelInputExample=None,
              **kwargs):
    Model.log(artifact_path=artifact_path, flavor=mlflow.xgboost,
              registered_model_name=registered_model_name,
              xgb_model=xgb_model, conda_env=conda_env,
              signature=signature, input_example=input_example, **kwargs)

mlflow.models.Modelのlogメソッド（classmethod）のラッパです。大まかには、以下の3ステップを実行します。

MLflow Modelの生成
artifactの生成
生成したMLflow Model（artifactを含む）をrunとして記録

なお、モデルのartifact化に使用するパラメータはModel.logのkwargs引数に入ってからflavor.save_model()呼び出し時に渡されます。 save_model()のxgb_model変数はkwargsを通じてsave_model()に渡され、artifact化されます。

load_model()

def load_model(model_uri):
    local_model_path = _download_artifact_from_uri(artifact_uri=model_uri)
    flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME)
    xgb_model_file_path = os.path.join(local_model_path, flavor_conf.get("data", "model.xgb"))
    return _load_model(path=xgb_model_file_path)

この関数の処理内容ですが、

artifactをダウンロードして
flavorの設定を読み込んで
xgboost.Boosterをartifactからロード

と、中で呼び出されている関数名を読んだ通りの処理内容です。

特に解説することもないのですが、 _download_artifact_from_uri()や_get_flavor_configuration()といった便利なutility関数が実装されていることを知っているとflavor実装などで役に立ちます。

おわりに

駆け足となりましたが、XGBoost flavorの内容とその前提となる知識を軽く紹介しました。シンプルなflavorなので、MLflowの理解のために読むのはもちろんですが、flavorを作ってみる際の手本としても最適だと思います。実際、今回書いていてすごく理解が深まり~~、前回記事の内容が思いっきり間違っていることにも気付け~~ました。

当該箇所は修正済みです。本当にすみませんでした……。

他にも、純粋にPythonモジュールの実装としても、多層に分けて段階的に抽象度を下げていくなど勉強になることが多かったです。汎化と実装コストのトレードオフにはいつも苦心しているので、今後書くプログラムに色々と取り入れようと思っています。

では、良きPython Lifeを！