MLFlow Pyfuncのいいところを理解したい

MLflow Pyfuncとは？

MLflowを使ってみると格納したモデルに対して、以下のようなガイドが現れます。このmlflow pyfuncとは何なのか？、何が嬉しいのか今回は掘り下げたいと思います。

先に結論

• 書き方を統一できるよ
• Spark使ってたら並列推論できるよ

まずは実際に利用してみよう

訓練用のcode

import mlflow
from sklearn.datasets import load_iris
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split

mlflow.set_tracking_uri("http://localhost:8080")

data = load_iris()
X = data.data
y = data.target

X_train, _, y_train, _ = train_test_split(X, y, test_size=0.2, random_state=42)

with mlflow.start_run(experiment_id="0"):
    model = LinearRegression()
    model.fit(X_train, y_train)
    mlflow.sklearn.log_model(model, "model")

推論用のcode

mlflow.set_tracking_uri("http://localhost:8080")

data = load_iris()
X = data.data
y = data.target

_, X_test, _, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

logged_model = 'runs:/7f1059605e29410eb2d15bc16fdaa9f4/model'
loaded_model = mlflow.pyfunc.load_model(logged_model)
result = loaded_model.predict(y_train)
print(result)

これで無事推論が実行され、下記のような結果を得られます。

array=[0 0 1 0 0 2 1 0 0 0 2 1 1 0 0 1 2 2 1 2 1 2 1 0 2 1 0 0 0 1 2 0 0 0 1 0 1
 2 0 1 2 0 2 2 1 1 2 1 0 1 2 0 0 1 1 0 2 0 0 1 1 2 1 2 2 1 0 0 2 2 0 0 0 1
 2 0 2 2 0 1 1 2 1 2 0 2 1 2 1 1 1 0 1 1 0 1 2 2 0 1 2 2 0 2 0 1 2 2 1 2 1
 1 2 2 0 1 2 0 1 2]

読んでみよう Understanding PyFunc in MLflow

一番の強みは？： 規格化と標準化

もう少し付け加えていうと、ニッチなML Frameworkに対して、全て同一のインターフェイスで扱えるようにしていますよ、ということである。

まず最初の一文が気になる。

while named flavors offer specific functionalities tailored to popular frameworks

named flavorsは人気であるが、というところから始まる。named flavorsとは、参照リンクは無かったが以下の内容で間違いない。
https://mlflow.org/docs/latest/models.html#built-in-model-flavors
中にはPytorch,keras,lightgbm,Transformersと、大概なんでも揃っている。
(逆にいうと、ここに自身の利用したいライブラリが存在しているのであれば、それを使えば十分とも読み取れる。)

しかし、これらが存在しないニッチなframeworkは独自のコーディングが必要とされ、全体の統一に支障がでるため、 MLflow pyfuncがそれらを統一すること出来る。ということがまず強みのようだ。
確かに複数モデルを使いアンサンブルなどしたい時には重宝しそうである。

大規模データに対する並列処理

これは Conditions where a Custom Pyfunc might be best
の項に書かれていた一つ目のメリットではあるが、かなり強いポイントだと感じる。
Spark等の並列分散処理が可能なエンジンを利用している場合、大規模データに対して並列推論できると書いてある。

Custom pyfuncによる拡張性

pyfuncのデメリット

この話をする前に一つデモを共有したい。

# sklearn flavor pattern
loaded_model = mlflow.sklearn.load_model(logged_model)
print('coefficient = ', loaded_model.coef_[0])
print('intercept = ', loaded_model.intercept_)

# pyfunc flavor pattern
loaded_model = mlflow.pyfunc.load_model(logged_model)
print('coefficient = ', loaded_model.coef_[0])
print('intercept = ', loaded_model.intercept_)

先ほどの推論コードを書き換え、上記のように会期モデルの傾きと切片を標準出力するコードに書き換えるとどうなるか？下記がそのoutputである。

coefficient =  -0.11633479416518296
intercept =  0.25252758981814805
Traceback (most recent call last):
  File "iris_inference_model.py", line 26, in <module>
    print('coefficient = ', loaded_model.coef_[0])
AttributeError: 'PyFuncModel' object has no attribute 'coef_'

１回目のprintは上手くいき、２つ目のprintは失敗する。
これはpyfuncで読み込まれたモデルには本来の sklearn.linear_model.LinearRegressionが持つメソッドを使えなくなることを示している。
pyfuncは規格化された別のクラスなので、当然そりゃそうと言えばそうなのですが、中身をsklearn.linear_model.LinearRegressionだからと言ってそれらを勝手には引き継いでくれないのです。

また、もう一つ見て欲しいnotebookもあり、これは公式が提示していたのですが、こういった面倒もあります。
https://mlflow.org/docs/latest/traditional-ml/creating-custom-pyfunc/notebooks/override-predict.html#Defining-our-custom-PythonModel
このnotebookはLogisticRegressionを使った時にpredict_probaが使えないというnotebookです。こちらも一度直面し、結構面を食らいました。

しかし Custom pyfuncが可能

という話なのですが、、、これは半分解決可能で、半分は解決できない認識です。
まず前節の最後に取り上げたnotebookの通り、predict_probaについては解決可能です。
しかし、 .predict(x_test, params={"predict_method": "predict_log_proba"}) と引数で管理することになり、少しめんどうですね。。。

もう一方coef_を取得する方法については、色々試してできませんでした。
（できた方いたら教えてください。）

例えば下記のコード、predictの中に潜ませたcoef_はprintされたものの、新しく作ってみたcoef_ methodは呼び出せませんでした。

from mlflow.pyfunc import PythonModel
from joblib import dump

class PyFuncLinearWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.model = None

    def load_context(self, context):
        from joblib import load
        self.model = load(context.artifacts["model_path"])
    
    def coef_(self):
        # こちらの勝手に作ったmeshodは利用できなかった。
        return self.model.coef_

    def predict(self, context, model_input, params=None):
        print(self.model.coef_)  # こんな感じでprintさせることは成功
        return self.model.predict(model_input)


with mlflow.start_run(experiment_id="0") as run:
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    model_directory = "/tmp/sklearn_model.joblib"
    dump(model, model_directory)
    artifacts = {"model_path": model_directory}
    mlflow.pyfunc.log_model(artifact_path="model", python_model=PyFuncLinearWrapper(), artifacts=artifacts, pip_requirements=["joblib", "sklearn"],)

この辺り結構コード量も複雑で、学習コストも低く無いので、便利ではあるものの、
sklearn等利用しているのであれば、普通にmlflow.sklearn.log_model とmlflow.sklearn.load_modelで十分対応できそうな気がします。

まとめ

• mlflow pyfuncはあらゆる ML frameworkの推論等を同じコードで書けるようにしてくれる。
• Spark等利用時には並列推論を実現できる。
• Wrapper等記述することが可能で、predict_proba等にも（wrapperを書けば・・・）対応可能である。

個人的にはメジャーなFrameworkを使うことがほとんどなので、これは使わないかも。