自定义模型的 predict 方法

下载此笔记本

在本教程中，我们将探讨在 MLflow 的 PyFunc flavor 中自定义模型的 predict 方法的过程。 当您希望在通过 MLflow 部署模型后，对其行为方式有更大的灵活性时，这将特别有用。

为了说明这一点，我们将使用著名的 Iris 数据集，并使用 scikit-learn 构建一个基本的 Logistic Regression 模型。

from joblib import dump
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split

import mlflow
from mlflow.models import infer_signature
from mlflow.pyfunc import PythonModel

配置跟踪服务器 URI

此步骤很重要，可以确保我们在此笔记本中进行的所有 MLflow 调用实际上都将记录到本地运行的跟踪服务器。

如果您在不同的环境中按照此笔记本进行操作，并希望将此笔记本的其余部分执行到远程跟踪服务器，请更改以下单元格。

Databricks：mlflow.set_tracking_uri("databricks")

您的托管 MLflow：mlflow.set_tracking_uri("http://my.company.mlflow.tracking.server:<port>)

您的本地跟踪服务器。 如入门教程中所述，我们可以通过命令行启动本地跟踪服务器，如下所示

mlflow server --host 127.0.0.1 --port 8080

MLflow UI 服务器可以通过以下方式在本地启动

mlflow ui --host 127.0.0.1 --port 8090

mlflow.set_tracking_uri("https://:8080")

让我们首先加载 Iris 数据集并将其拆分为训练集和测试集。 然后，我们将在训练数据上训练一个简单的 Logistic Regression 模型。

iris = load_iris()
x = iris.data[:, 2:]
y = iris.target

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=9001)

model = LogisticRegression(random_state=0, max_iter=5_000, solver="newton-cg").fit(x_train, y_train)

这是机器学习中常见的场景。 我们有一个经过训练的模型，我们想用它来进行预测。 使用 scikit-learn，该模型提供了几种方法来实现此目的

• predict - 用于预测类标签
• predict_proba - 用于获取类成员概率
• predict_log_proba - 用于获取每个类的对数概率

我们可以预测类标签，如下所示。

model.predict(x_test)[:5]

array([1, 2, 2, 1, 0])

我们还可以获得类成员资格概率。

model.predict_proba(x_test)[:5]

array([[2.64002987e-03, 6.62306827e-01, 3.35053144e-01],
     [1.24429110e-04, 8.35485037e-02, 9.16327067e-01],
     [1.30646549e-04, 1.37480519e-01, 8.62388835e-01],
     [3.70944840e-03, 7.13202611e-01, 2.83087941e-01],
     [9.82629868e-01, 1.73700532e-02, 7.88350143e-08]])

以及生成每个类的对数概率。

model.predict_log_proba(x_test)[:5]

array([[ -5.93696505,  -0.41202635,  -1.09346612],
     [ -8.99177441,  -2.48232793,  -0.08738192],
     [ -8.94301498,  -1.98427305,  -0.14804903],
     [ -5.59687209,  -0.33798973,  -1.26199768],
     [ -0.01752276,  -4.05300763, -16.35590859]])

虽然在同一 Python 会话中直接使用该模型很简单，但是当我们想保存此模型并在其他地方加载它时会发生什么，尤其是在使用 MLflow 的 PyFunc flavor 时？ 让我们探索这种情况。

mlflow.set_experiment("Overriding Predict Tutorial")

sklearn_path = "/tmp/sklearn_model"

with mlflow.start_run() as run:
  mlflow.sklearn.save_model(
      sk_model=model,
      path=sklearn_path,
      input_example=x_train[:2],
  )

/Users/benjamin.wilson/miniconda3/envs/mlflow-dev-env/lib/python3.8/site-packages/_distutils_hack/__init__.py:30: UserWarning: Setuptools is replacing distutils.
warnings.warn("Setuptools is replacing distutils.")

一旦该模型作为 pyfunc 加载，默认行为仅支持 predict 方法。 当您尝试调用其他方法（如 predict_proba）时，这一点很明显，这会导致 AttributeError。 这可能会受到限制，尤其是在您想保留原始模型的全部功能时。

loaded_logreg_model = mlflow.pyfunc.load_model(sklearn_path)

loaded_logreg_model.predict(x_test)

array([1, 2, 2, 1, 0, 1, 2, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 1, 1, 2, 1,
     1, 0, 1, 1, 0, 0, 1, 2])

这完全按照我们的预期工作。 输出与保存之前直接使用模型时的输出相同。

让我们尝试使用 predict_proba 方法。

我们实际上不会运行它，因为它会引发异常。 这是我们尝试执行此操作时的行为

loaded_logreg_model.predict_proba(x_text)

这将导致此错误

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/var/folders/cd/n8n0rm2x53l_s0xv_j_xklb00000gp/T/ipykernel_15410/1677830262.py in <cell line: 1>()
----> 1 loaded_logreg_model.predict_proba(x_text)

AttributeError: 'PyFuncModel' object has no attribute 'predict_proba'

部署后，我们该如何支持模型的原始行为？

我们可以创建一个自定义 pyfunc，该 pyfunc 会覆盖 predict 方法的行为。

对于下面的示例，我们将展示 pyfunc 的两个功能，这些功能可用于处理自定义模型日志记录功能

• 覆盖 predict 方法
• 自定义加载工件

需要注意的关键一点是使用 joblib 进行序列化。 虽然 pickle 历史上一直用于序列化 scikit-learn 模型，但现在建议使用 joblib，因为它提供了更好的性能和支持，尤其是在大型 numpy 数组方面。

我们将使用 joblib 及其 dump 和 load API 来处理将我们的模型对象加载到我们的自定义 pyfunc 实现中。 使用 load_context 方法来处理加载文件（在实例化 pyfunc 对象时）的过程对于具有非常大或众多工件依赖项的模型（例如 LLM）特别有用，并且可以帮助显着减少在分布式系统中（例如 Apache Spark 或 Ray）加载的 pyfunc 的总内存占用空间。

from joblib import dump

from mlflow.models import infer_signature
from mlflow.pyfunc import PythonModel

要了解如何在自定义 Python 模型中利用 load_context 功能，我们将首先使用 joblib 在本地序列化我们的模型。 此处使用 joblib 纯粹是为了演示一种非标准方法（一种 MLflow 本机不支持的方法），以说明 Python 模型实现的灵活性。 只要我们在 load_context 中导入此库，并且在我们将加载此模型的环境中可用，模型工件将被正确地反序列化。

model_directory = "/tmp/sklearn_model.joblib"
dump(model, model_directory)

['/tmp/sklearn_model.joblib']

定义我们的自定义 PythonModel

下面的 ModelWrapper 类是自定义 pyfunc 的一个示例，它扩展了 MLflow 的 PythonModel。 它通过使用 predict method 的 params 参数来提供预测方法的灵活性。 这样，我们就可以在调用加载的 pyfunc 实例上的 predict 方法时，指定我们想要常规的 predict、predict_proba 还是 predict_log_proba 行为。

class ModelWrapper(PythonModel):
  def __init__(self):
      self.model = None

  def load_context(self, context):
      from joblib import load

      self.model = load(context.artifacts["model_path"])

  def predict(self, context, model_input, params=None):
      params = params or {"predict_method": "predict"}
      predict_method = params.get("predict_method")

      if predict_method == "predict":
          return self.model.predict(model_input)
      elif predict_method == "predict_proba":
          return self.model.predict_proba(model_input)
      elif predict_method == "predict_log_proba":
          return self.model.predict_log_proba(model_input)
      else:
          raise ValueError(f"The prediction method '{predict_method}' is not supported.")

在定义自定义 pyfunc 之后，后续步骤包括使用 MLflow 保存模型，然后再将其加载回来。 加载的模型将保留我们构建到自定义 pyfunc 中的灵活性，从而允许我们动态选择预测方法。

注意：下面的 artifacts 引用非常重要。 为了使 load_context 能够访问我们指定为已保存模型位置的路径，必须将其作为字典提供，该字典将适当的访问键映射到相关值。 如果未将此字典作为 mlflow.save_model() 或 mlflow.log_model() 的一部分提供，则将导致此自定义 pyfunc 模型无法正确加载。

# Define the required artifacts associated with the saved custom pyfunc
artifacts = {"model_path": model_directory}

# Define the signature associated with the model
signature = infer_signature(x_train, params={"predict_method": "predict_proba"})

我们可以看到已定义的参数如何在签名定义中使用。 如下所示，参数在记录时会发生细微的变化。 我们有一个已定义的参数键 (predict_method)、预期类型 (string) 和默认值。 这最终意味着此 params 定义为

• 我们只能为键 predict_method 提供 params 覆盖。 除此之外的任何内容都将被忽略，并且将显示一条警告，指示未知的参数将不会传递给底层模型。

• 与 predict_method 关联的值必须是字符串。 任何其他类型都是不允许的，并且会引发意外类型的异常。

• 如果在调用 predict 时未提供 predict_method 的任何值，则该模型将使用 predict_proba 的默认值。

signature

inputs: 
[Tensor('float64', (-1, 2))]
outputs: 
None
params: 
['predict_method': string (default: predict_proba)]

我们现在可以保存我们的自定义模型。 我们正在提供一个路径来保存它，以及包含我们通过 joblib 手动序列化的实例的位置的 artifacts 定义。 还包括 signature，这是使此示例正常工作的关键组件； 如果没有在签名中定义的参数，我们将无法覆盖 predict 方法将使用的预测方法。

注意，我们在这里覆盖了 pip_requirements，以确保我们指定两个依赖库的要求：joblib 和 sklearn。 这有助于确保我们将此模型部署到的任何环境都会在加载此已保存模型之前预加载这两个依赖项。

pyfunc_path = "/tmp/dynamic_regressor"

with mlflow.start_run() as run:
  mlflow.pyfunc.save_model(
      path=pyfunc_path,
      python_model=ModelWrapper(),
      input_example=x_train,
      signature=signature,
      artifacts=artifacts,
      pip_requirements=["joblib", "sklearn"],
  )

Downloading artifacts:   0%|          | 0/1 [00:00<?, ?it/s]

我们现在可以使用 mlflow.pyfunc.load_model API 重新加载我们的模型。

loaded_dynamic = mlflow.pyfunc.load_model(pyfunc_path)

让我们看看 pyfunc 模型在没有覆盖 params 参数的情况下会产生什么。

loaded_dynamic.predict(x_test)

array([[2.64002987e-03, 6.62306827e-01, 3.35053144e-01],
     [1.24429110e-04, 8.35485037e-02, 9.16327067e-01],
     [1.30646549e-04, 1.37480519e-01, 8.62388835e-01],
     [3.70944840e-03, 7.13202611e-01, 2.83087941e-01],
     [9.82629868e-01, 1.73700532e-02, 7.88350143e-08],
     [6.54171552e-03, 7.54211950e-01, 2.39246334e-01],
     [2.29127680e-06, 1.29261337e-02, 9.87071575e-01],
     [9.71364952e-01, 2.86348857e-02, 1.62618524e-07],
     [3.36988442e-01, 6.61070371e-01, 1.94118691e-03],
     [9.81908726e-01, 1.80911360e-02, 1.38374097e-07],
     [9.70783357e-01, 2.92164276e-02, 2.15395762e-07],
     [6.54171552e-03, 7.54211950e-01, 2.39246334e-01],
     [1.06968794e-02, 8.88253152e-01, 1.01049969e-01],
     [3.35084116e-03, 6.57732340e-01, 3.38916818e-01],
     [9.82272901e-01, 1.77269948e-02, 1.04445227e-07],
     [9.82629868e-01, 1.73700532e-02, 7.88350143e-08],
     [1.62626101e-03, 5.43474542e-01, 4.54899197e-01],
     [9.82629868e-01, 1.73700532e-02, 7.88350143e-08],
     [5.55685308e-03, 8.02036140e-01, 1.92407007e-01],
     [1.01733783e-02, 8.62455340e-01, 1.27371282e-01],
     [1.43317140e-08, 1.15653085e-03, 9.98843455e-01],
     [4.33536629e-02, 9.32351526e-01, 2.42948113e-02],
     [3.97007654e-02, 9.08506559e-01, 5.17926758e-02],
     [9.19762712e-01, 8.02357267e-02, 1.56085268e-06],
     [4.21970838e-02, 9.26463030e-01, 3.13398863e-02],
     [3.13635521e-02, 9.17295925e-01, 5.13405229e-02],
     [9.77454643e-01, 2.25452265e-02, 1.30412321e-07],
     [9.71364952e-01, 2.86348857e-02, 1.62618524e-07],
     [3.23802803e-02, 9.27626313e-01, 3.99934070e-02],
     [1.21876019e-06, 1.79695714e-02, 9.82029210e-01]])

正如预期的那样，它返回了 params predict_method 的默认值，即 predict_proba。 我们现在可以尝试覆盖该功能以返回类预测。

loaded_dynamic.predict(x_test, params={"predict_method": "predict"})

array([1, 2, 2, 1, 0, 1, 2, 0, 1, 0, 0, 1, 1, 1, 0, 0, 1, 0, 1, 1, 2, 1,
     1, 0, 1, 1, 0, 0, 1, 2])

我们也可以覆盖它以返回类成员资格的 predict_log_proba 对数概率。

loaded_dynamic.predict(x_test, params={"predict_method": "predict_log_proba"})

array([[-5.93696505e+00, -4.12026346e-01, -1.09346612e+00],
     [-8.99177441e+00, -2.48232793e+00, -8.73819177e-02],
     [-8.94301498e+00, -1.98427305e+00, -1.48049026e-01],
     [-5.59687209e+00, -3.37989732e-01, -1.26199768e+00],
     [-1.75227629e-02, -4.05300763e+00, -1.63559086e+01],
     [-5.02955584e+00, -2.82081850e-01, -1.43026157e+00],
     [-1.29864013e+01, -4.34850415e+00, -1.30127244e-02],
     [-2.90530299e-02, -3.55312953e+00, -1.56318587e+01],
     [-1.08770665e+00, -4.13894984e-01, -6.24445569e+00],
     [-1.82569224e-02, -4.01233318e+00, -1.57933050e+01],
     [-2.96519488e-02, -3.53302414e+00, -1.53507887e+01],
     [-5.02955584e+00, -2.82081850e-01, -1.43026157e+00],
     [-4.53780322e+00, -1.18498496e-01, -2.29214015e+00],
     [-5.69854387e+00, -4.18957208e-01, -1.08200058e+00],
     [-1.78861062e-02, -4.03266667e+00, -1.60746030e+01],
     [-1.75227629e-02, -4.05300763e+00, -1.63559086e+01],
     [-6.42147176e+00, -6.09772414e-01, -7.87679430e-01],
     [-1.75227629e-02, -4.05300763e+00, -1.63559086e+01],
     [-5.19272332e+00, -2.20601610e-01, -1.64814232e+00],
     [-4.58798095e+00, -1.47971911e-01, -2.06064898e+00],
     [-1.80607910e+01, -6.76233040e+00, -1.15721450e-03],
     [-3.13836408e+00, -7.00453618e-02, -3.71749248e+00],
     [-3.22638481e+00, -9.59531718e-02, -2.96050653e+00],
     [-8.36395634e-02, -2.52278639e+00, -1.33702783e+01],
     [-3.16540417e+00, -7.63811370e-02, -3.46286367e+00],
     [-3.46210882e+00, -8.63251488e-02, -2.96927492e+00],
     [-2.28033892e-02, -3.79223192e+00, -1.58525647e+01],
     [-2.90530299e-02, -3.55312953e+00, -1.56318587e+01],
     [-3.43020568e+00, -7.51263075e-02, -3.21904066e+00],
     [-1.36176765e+01, -4.01907543e+00, -1.81342258e-02]])

我们已经成功创建了一个 pyfunc 模型，该模型保留了原始 scikit-learn 模型的全部功能，同时使用了一种避免标准 pickle 方法的自定义加载器方法。

本教程重点介绍了 MLflow 的 PyFunc flavor 的强大功能和灵活性，演示了如何对其进行定制以满足您的特定需求。 在您继续构建和部署模型时，请考虑如何使用自定义 pyfunc 来增强模型的功能并适应各种情况。