创建自定义模型："Add N" 模型

下载此笔记本 我们的第一个示例简单但具有说明性。我们将创建一个模型，将指定的数值 n 添加到 Pandas DataFrame 输入的所有列中。这将演示定义自定义模型、保存模型、重新加载模型以及执行预测的过程。

步骤 1：定义模型类

我们首先为我们的模型定义一个 Python 类。这个类应该继承自 mlflow.pyfunc.PythonModel 并实现必要的方法。

import mlflow.pyfunc


class AddN(mlflow.pyfunc.PythonModel):
  """
  A custom model that adds a specified value `n` to all columns of the input DataFrame.

  Attributes:
  -----------
  n : int
      The value to add to input columns.
  """

  def __init__(self, n):
      """
      Constructor method. Initializes the model with the specified value `n`.

      Parameters:
      -----------
      n : int
          The value to add to input columns.
      """
      self.n = n

  def predict(self, context, model_input, params=None):
      """
      Prediction method for the custom model.

      Parameters:
      -----------
      context : Any
          Ignored in this example. It's a placeholder for additional data or utility methods.

      model_input : pd.DataFrame
          The input DataFrame to which `n` should be added.

      params : dict, optional
          Additional prediction parameters. Ignored in this example.

      Returns:
      --------
      pd.DataFrame
          The input DataFrame with `n` added to all columns.
      """
      return model_input.apply(lambda column: column + self.n)

步骤 2：保存模型

定义好模型类后，我们可以实例化它并使用 MLflow 进行保存。

# Define the path to save the model
model_path = "/tmp/add_n_model"

# Create an instance of the model with `n=5`
add5_model = AddN(n=5)

# Save the model using MLflow
mlflow.pyfunc.save_model(path=model_path, python_model=add5_model)

/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：加载模型

模型保存后，我们可以使用 MLflow 将其重新加载，然后用于预测。

# Load the saved model
loaded_model = mlflow.pyfunc.load_model(model_path)

步骤 4：评估模型

现在，让我们使用加载的模型对样本输入执行预测并验证其正确性。

import pandas as pd

# Define a sample input DataFrame
model_input = pd.DataFrame([range(10)])

# Use the loaded model to make predictions
model_output = loaded_model.predict(model_input)

model_output

	0	1	2	3	4	5	6	7	8	9
0	5	6	7	8	9	10	11	12	13	14

结论

这个简单的例子展示了 MLflow 自定义 pyfunc 的强大功能和灵活性。通过封装任意 Python 代码及其依赖项，自定义 pyfunc 模型为各种用例提供了统一且一致的接口。无论你是使用小众的机器学习框架，需要自定义预处理步骤，还是想要集成独特的预测逻辑，pyfunc 都是完成任务的工具。