跳到主要内容

MLflow Custom Pyfunc 入门

下载此笔记本

在本 MLflow Custom Pyfunc 入门教程中，我们将深入探讨 PythonModel 类的核心功能，并探索如何使用它们来构建一个可以保存、加载和用于推理的非常简单的模型。

目标：在本指南结束时，您将学会如何

  • 使用 Python 类定义自定义 PyFunc 模型。
  • 理解 PyFunc flavor 的核心组件。
  • 保存、加载和预测使用自定义 PyFunc 模型。
  • 利用 MLflow PyFunc 的强大功能实现一个实际示例：利萨如图形。

PythonModel 类​

MLflow 对于通用模型实例类型的处理方式是严格标准化的，以确保任何存储在 MLflow 中的模型都可以用于推理，前提是遵循实现指南。

有两种方法可以创建自定义 PythonModel 实例。第一种方法，也是本指南将要使用的方法，涉及定义一个类和将用于接口的方法。还有另一种方法，即定义一个名为 predict 的函数，并将其作为 mlflow.pyfunc.save_model() 中的 python_model 参数进行记录。这种方法功能有限，但对于可以将整个预测逻辑封装在单个函数中的实现来说，它是更可取的。对于第二种 pyfunc 记录模式，MLflow 将为您创建一个通用的 PythonModel 类并进行记录，并将您提供的 predict 函数添加为类中的 predict() 方法。

核心 PythonModel 组件​

MLflow 的 PyFunc 围绕 PythonModel 类展开。该类中的两个基本方法是

  1. load_context(self, context)：此方法用于加载工件或其他初始化任务。它是可选的，可用于获取外部引用。
  2. predict(self, context, model_input, params=None)：这是您的模型在进行预测时的入口点。您的自定义 PyFunc 模型必须定义此方法。

例如，如果您的模型使用 XGBoost 等外部库，您可以在 load_context 方法中加载 XGBoost 模型，并在 predict 方法中使用它。

PythonModel 的基本指南​

此方法的指南如下：

  1. 您的类必须是 mlflow.pyfunc.PythonModel 的子类。
  2. 您的类必须实现 predict 方法。
  3. predict 方法必须符合 推理 API 的要求。
  4. predict 方法必须将 context 作为第一个命名参数引用。
  5. 如果您希望通过模型提供参数，则必须将其定义为模型签名的一部分。签名必须与模型一起保存。
  6. 如果您打算在加载模型时执行其他功能（例如加载其他依赖文件），则可以决定在类中定义 load_context 方法。

定义一个简单的 Python 模型​

在本教程中，我们不会介绍更高级的 load_context 或在 predict 方法中与 context 参数的交互。为了理解自定义 PythonModel 的最基本方面，我们将保持简单。

为了展示 MLflow 自定义 Pyfunc 模型的一些额外用法，我们不会研究典型的库用例。相反，我们将研究如何使用 MLflow 存储利萨如图形实现的配置实例。

利萨如图形​

利萨如图形起源于谐波领域，是参数化正弦曲线，定义为：

$$ x(t) = A \sin(a t + \delta) $$ $$ y(t) = B \sin(b t) $$

其中

  • ( $A$ ) 和 ( $B$ ) 分别是 x 轴和 y 轴上曲线的振幅。

  • ( $a$ ) 和 ( $b$ ) 决定了振荡的频率。

  • ( $\delta$ ) 是 x 和 y 分量之间的相位差。

我们将创建一个简单的模型，允许用户生成与频率振荡比率及其相位相关的不同图案。

步骤 1：定义自定义 PyFunc 模型​

我们首先为自定义模型定义一个 Python 类。该类应继承自 mlflow.pyfunc.PythonModel。

在我们的利萨如图形模型中，我们使用参数 ( $A$ )、( $B$ ) 和 num_points 进行初始化。predict 方法负责根据输入 ( $a$ )、( $b$ ) 和 ( $\delta$ ) 绘制利萨如图形。

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

import mlflow.pyfunc
from mlflow.models import infer_signature


class Lissajous(mlflow.pyfunc.PythonModel):
  def __init__(self, A=1, B=1, num_points=1000):
      self.A = A
      self.B = B
      self.num_points = num_points
      self.t_range = (0, 2 * np.pi)

  def generate_lissajous(self, a, b, delta):
      t = np.linspace(self.t_range[0], self.t_range[1], self.num_points)
      x = self.A * np.sin(a * t + delta)
      y = self.B * np.sin(b * t)
      return pd.DataFrame({"x": x, "y": y})

  def predict(self, context, model_input, params=None):
      """
      Generate and plot the Lissajous curve with annotations for parameters.

      Args:
      - model_input (pd.DataFrame): DataFrame containing columns 'a' and 'b'.
      - params (dict, optional): Dictionary containing optional parameter 'delta'.
      """
      # Extract a and b values from the input DataFrame
      a = model_input["a"].iloc[0]
      b = model_input["b"].iloc[0]

      # Extract delta from params or set it to 0 if not provided
      delta = params.get("delta", 0)

      # Generate the Lissajous curve data
      df = self.generate_lissajous(a, b, delta)

      sns.set_theme()

      # Create the plot components
      fig, ax = plt.subplots(figsize=(10, 8))
      ax.plot(df["x"], df["y"])
      ax.set_title("Lissajous Curve")

      # Define the annotation string
      annotation_text = f"""
      A = {self.A}
      B = {self.B}
      a = {a}
      b = {b}
      delta = {np.round(delta, 2)} rad
      """

      # Add the annotation with a bounding box outside the plot area
      ax.annotate(
          annotation_text,
          xy=(1.05, 0.5),
          xycoords="axes fraction",
          fontsize=12,
          bbox={"boxstyle": "round,pad=0.25", "facecolor": "aliceblue", "edgecolor": "black"},
      )

      # Adjust plot borders to make space for the annotation
      plt.subplots_adjust(right=0.65)
      plt.close()

      # Return the plot
      return fig

步骤 2：保存模型​

定义模型类后，我们可以对其进行实例化并使用 MLflow 进行保存。infer_signature 方法在这里很有用，可以自动推断模型的输入和输出模式。

由于我们使用 params 来覆盖方程式的 delta 值，因此我们需要在保存时提供模型的签名。如果此处未定义签名，则加载此模型实例的使用（如果保存时没有签名）将忽略提供的 params（并发出警告）。

# Define the path to save the model
model_path = "lissajous_model"

# Create an instance of the model, overriding the default instance variables `A`, `B`, and `num_points`
model_10k_standard = Lissajous(1, 1, 10_000)

# Infer the model signature, ensuring that we define the params that will be available for customization at inference time
signature = infer_signature(pd.DataFrame([{"a": 1, "b": 2}]), params={"delta": np.pi / 5})

# Save our custom model to the path we defined, with the signature that we declared
mlflow.pyfunc.save_model(path=model_path, python_model=model_10k_standard, signature=signature)
/Users/benjamin.wilson/miniconda3/envs/mlflow-dev-env/lib/python3.8/site-packages/mlflow/models/signature.py:212: UserWarning: Hint: Inferred schema contains integer column(s). Integer columns in Python cannot represent missing values. If your input data contains missing values at inference time, it will be encoded as floats and will cause a schema enforcement error. The best way to avoid this problem is to infer the model schema based on a realistic data sample (training dataset) that includes missing values. Alternatively, you can declare integer columns as doubles (float64) whenever these columns may have missing values. See `Handling Integers With Missing Values <https://www.mlflow.org/docs/latest/models.html#handling-integers-with-missing-values>`_ for more details.
inputs = _infer_schema(model_input) if model_input is not None else None
/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.")

步骤 3：加载模型​

保存后，我们可以将模型加载回来并用于预测。在这里，我们的预测是利萨如图形绘图。

# Load our custom model from the local artifact store
loaded_pyfunc_model = mlflow.pyfunc.load_model(model_path)

步骤 4：使用模型生成图形​

# Define the input DataFrame. In our custom model, we're reading only the first row of data to generate a plot.
model_input = pd.DataFrame({"a": [3], "b": [2]})

# Define a params override for the `delta` parameter
params = {"delta": np.pi / 3}

# Run predict, which will call our internal method `generate_lissajous` before generating a `matplotlib` plot showing the curve
fig = loaded_pyfunc_model.predict(model_input, params)

# Display the plot
fig

# Try a different configuration of arguments
fig2 = loaded_pyfunc_model.predict(
  pd.DataFrame({"a": [15], "b": [17]}), params={"delta": np.pi / 5}
)

fig2