使用 mlflow 提供自定义 python 评分模型

Use mlflow to serve a custom python model for scoring

我正在使用带有 mlflow 的 ml 软件生成的 Python 代码来读取数据帧,执行一些 table 操作并输出数据帧。我能够 运行 代码成功并将新数据框保存为工件。但是我无法使用 log_model 记录模型,因为它不是我们训练和拟合的 lr 或分类器模型。我想为此记录一个模型,以便可以为它提供新数据并进行部署 API

df = pd.read_csv(r"/home/xxxx.csv")


with mlflow.start_run():
    def getPrediction(row):
        
        perform_some_python_operaions 

        return [Status_prediction, Status_0_probability, Status_1_probability]
    columnValues = []
    for column in columns:
        columnValues.append([])

    for index, row in df.iterrows():
        results = getPrediction(row)
        for n in range(len(results)):
            columnValues[n].append(results[n])

    for n in range(len(columns)):
        df[columns[n]] = columnValues[n]

    df.to_csv('dataset_statistics.csv')
    mlflow.log_artifact('dataset_statistics.csv')
   

MLflow 支持 custom models 的 mlflow.pyfunc 风格。您可以创建一个继承自 mlflow.pyfunc.PythonModel 的自定义 class,它需要提供用于执行预测的函数 predict,以及可选的 load_context 来加载必要的工件,就像这样(采用来自文档):

class MyModel(mlflow.pyfunc.PythonModel):

    def load_context(self, context):
        # load your artifacts

    def predict(self, context, model_input):
        return my_predict(model_input.values)

您可以将模型所需的任何工件记录到 MLflow,必要时定义 Conda 环境等。
然后你可以使用 save_model 和你的 class 来保存你的实现,可以用 load_model 加载并使用你的模型执行 predict:

mlflow.pyfunc.save_model(
        path=mlflow_pyfunc_model_path, 
        python_model=MyModel(), 
        artifacts=artifacts)

# Load the model in `python_function` format
loaded_model = mlflow.pyfunc.load_model(mlflow_pyfunc_model_path)