如何从MLflow已注册模型中高效获取预处理参数？

问题背景 #

我正在学习MLflow，目前主要用它做实验追踪，未来打算把它作为集中式模型数据库，这样针对特定任务更新模型后，本地部署新版本时不用修改API。我的场景里，推理数据传入模型前需要预处理，而且这些预处理参数是模型开发的一部分，所以推理时必须获取这些参数来准备模型输入。

目前我把预处理参数以JSON格式附加到MLflow运行（run）里，但注册模型时这些参数并没有被包含进去。我现在是本地操作（通过UI注册选定的模型），但迁移到MLflow服务器后希望方案更稳健。我现在可以通过已注册模型的metadata.run_id获取对应的artifact，但有没有更优的方法？

当前实现代码：

model_URI = "models:/foo"
model = mlflow.pyfunc.load_model(model_URI)
runID = model.metadata.run_ID
params_path = "runs:/" + run_ID + "/params.json"
params = mlflow.artifacts.load_dict(params_path)

更优方案推荐 #

1. 封装预处理逻辑与参数为自定义PyFunc模型 #

将预处理参数和逻辑直接打包进模型，加载模型后可直接获取参数并执行预处理，无需额外读取run的artifact，是生产环境最推荐的方案：

import mlflow.pyfunc
import json

class PreprocessingWrappedModel(mlflow.pyfunc.PythonModel):
    def __init__(self, preprocessing_params, base_model):
        self.preprocessing_params = preprocessing_params
        self.base_model = base_model

    def predict(self, context, model_input):
        # 基于预处理参数执行数据预处理
        processed_input = self._preprocess(model_input)
        return self.base_model.predict(processed_input)

    def _preprocess(self, input_data):
        # 实现你的预处理逻辑，例如标准化、编码等
        # 示例：使用scaler_mean和scaler_std做标准化
        return (input_data - self.preprocessing_params["scaler_mean"]) / self.preprocessing_params["scaler_std"]

# 假设已训练好基础模型base_model，以及定义好预处理参数
preprocessing_params = {"scaler_mean": 0.5, "scaler_std": 0.2}

# 启动run并保存自定义模型
with mlflow.start_run() as run:
    mlflow.pyfunc.log_model(
        artifact_path="wrapped_model",
        python_model=PreprocessingWrappedModel(preprocessing_params, base_model),
        # 可选：将参数单独保存为artifact作为备份
        artifacts={"preprocessing_params": json.dumps(preprocessing_params)}
    )
    # 注册模型到模型仓库
    mlflow.register_model(f"runs:/{run.info.run_id}/wrapped_model", "foo")

加载模型时直接获取参数：

model = mlflow.pyfunc.load_model("models:/foo")
preprocessing_params = model._model_impl.preprocessing_params

2. 将预处理参数存入模型的Tags #

注册模型时，把预处理参数序列化为字符串存入模型的tags中，加载模型后可直接从metadata读取，适合参数内容较少的场景：

import json
from mlflow.tracking import MlflowClient

client = MlflowClient()
run_id = "你的运行ID"
preprocessing_params = {"scaler_mean": 0.5, "scaler_std": 0.2}

# 注册模型时添加自定义tags
model_version = client.create_model_version(
    name="foo",
    source=f"runs:/{run_id}/model",
    run_id=run_id,
    tags={"preprocessing_params": json.dumps(preprocessing_params)}
)

# 加载模型后获取参数
model = mlflow.pyfunc.load_model("models:/foo")
params_str = model.metadata.tags.get("preprocessing_params")
preprocessing_params = json.loads(params_str)

3. 将预处理参数与模型存放在同一Artifact目录 #

在训练run中，把params.json和模型文件放在同一个artifact目录下，注册模型时会自动包含该文件，加载模型后直接从模型本地路径读取：

import os
import json
import mlflow.sklearn

# 假设已训练好模型base_model和定义好预处理参数
preprocessing_params = {"scaler_mean": 0.5, "scaler_std": 0.2}

with mlflow.start_run() as run:
    # 保存模型到artifact目录
    mlflow.sklearn.log_model(base_model, artifact_path="model")
    # 将预处理参数保存到模型同一目录下
    mlflow.log_dict(preprocessing_params, "model/params.json")
    # 注册模型
    mlflow.register_model(f"runs:/{run.info.run_id}/model", "foo")

加载模型后读取参数：

import os
import json

model = mlflow.pyfunc.load_model("models:/foo")
# 获取模型在本地的存储路径
model_local_path = model.metadata._model_impl._model_path
params_path = os.path.join(model_local_path, "params.json")
with open(params_path, "r") as f:
    preprocessing_params = json.load(f)

方案对比 #

内容的提问来源于stack exchange，提问作者Aleksander Marek