创建自定义模型：“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 都是适合您的工具。