从PythonModelContext返回的“mlflow.pyfunc.load_model”对象--如何检索保存的原始模型

Question

我正在创建一个自定义myflow.pyfunc对象，希望将其保存到MLFlow中，然后再检索。我不明白用mlflow.pyfunc.save_model()保存的对象和用mlflow.pyfunc.load_model()检索的对象之间的关系。

加载的模型是一个'PythonModelContext‘对象，而不是我原来的python类。当我试图在加载的版本中使用预测方法时，我会得到一个错误。

在这里，我初始化了MLflow并创建了类的虚拟示例

# load 
import os
import tempfile
from pathlib import Path
import pandas as pd
import mlflow
from mlflow.tracking import MlflowClient
import mlflow.pyfunc
from mlflow.pyfunc import PythonModelContext

# initialise MLFlow
mlflow_var = os.getenv('HYMIND_REPO_TRACKING_URI')
mlflow.set_tracking_uri(mlflow_var)   

client = MlflowClient()

# Define the class that will be used for fit and predict (dummy example)
class PredictSpeciality(mlflow.pyfunc.PythonModel):
    
    def fit(self):
        print('fit')
        d = {'col1': [1, 2], 'col2': [3, 4]}
        df = pd.DataFrame(data=d)
        return df
           
    def predict(self, X, y=None):
        print('predict')
        print(X.shape)
        return 

如果我现在像运行类一样运行这个类，那么预测方法就能工作：

# Use of this predictor before saving works fine 
m = PredictSpeciality()
df = m.fit()
m.predict(df)

但是，如果我将模型保存到注册表，然后重新加载它，则预测方法将不再起作用：

counter +=1
exp_name = 'MLflow-test-' + str(counter)

os.environ["MLFLOW_EXPERIMENT_NAME"] = exp_name
experiment_id = mlflow.create_experiment(exp_name)

mlflow.set_experiment(exp_name)
experiment = dict(mlflow.get_experiment_by_name(exp_name))
experiment_id = experiment['experiment_id']

with mlflow.start_run():
    
    # dummy code here for fitting a model
    m = PredictSpeciality()
    df = m.fit()
    
# mark best run
runs = mlflow.search_runs()
best_run_id = runs['run_id'][0]

# tag the best run and save model
with mlflow.start_run(run_id=best_run_id):
    mlflow.set_tag('best_run_', 1)   

    mlflow_model_path = f'/data/hymind/repo/{experiment_id}/{best_run_id}/artifacts/model/'
    mlflow.pyfunc.save_model(path=mlflow_model_path, python_model=m)
    
# end experiment and register best model
model_name = 'MLflow-test' + str(counter)
registered_model = mlflow.register_model(f'runs:/{best_run_id}/model', model_name)

# now attempt to make a prediction using the loaded model
model_version = 1
m = mlflow.pyfunc.load_model(f"models:/{model_name}/{model_version}")
m.predict(df)

在这种情况下，我会得到属性错误。

AttributeError：'PythonModelContext‘对象没有属性'shape’

如何从'PythonModelContext‘对象中获取原始模型？

Answer 1

如果仔细查看正在扩展的predict()类中的抽象方法mlflow.pyfunc.PythonModel的签名，您将看到它有3个参数：

def predict(self, context, model_input):

因此，如果您将简单的类更改为具有额外的参数context，那么您的示例应该可以工作：

class PredictSpeciality(mlflow.pyfunc.PythonModel):
    def fit(self):
        print('fit')
        d = {'col1': [1, 2], 'col2': [3, 4]}
        df = pd.DataFrame(data=d)
        return df
           
    def predict(self, context, X, y=None):
        print('predict')
        print(X.shape)
        return 

更详细地介绍一下这里发生了什么:有两个类在起作用：mlflow.pyfunc.PythonModel和mlflow.pyfunc.PyFuncModel。

mlflow.pyfunc.PythonModel正在被mlflow.pyfunc.PyFuncModel包装。前者是实际工作，后者是处理元数据、打包、conda环境等。在文档中解释如下：

Python函数模型作为mlflow.pyfunc.PyFuncModel的一个实例加载，它是围绕模型实现和模型元数据(MLmodel文件)的MLflow包装器。

不幸的是，您不能直接创建一个文件还指出，但只能

模型实现和元数据的包装器。这个类不是直接构造的。相反，这个类的实例是从mlflow.pyfunc.load_model()构造和返回的。

我发现这是相当有限的，我不知道为什么它是这样设计的，但是，在这里你可以做两件事：

1. 在直接处理包装类时，传入一个额外的参数：

   m.predict(None, df)

1. 保存并加载模型以获得mlflow.pyfunc.PyFuncModel：

   mlflow.pyfunc.save_model(path="temp_model", python_model=m)
   m2 = mlflow.pyfunc.load_model("temp_model")
   m2.predict(df)

我知道这并不优雅，但实际上我过去一直在使用2号。如果MLFlow团队的人能够评论为什么不支持直接创建mlflow.pyfunc.PyFuncModel，那就太好了。