使用 mlflow.pyfunc 的自定义 MLflow 模型

如果您想了解 MLflow 自定义模型在灵活性和定制性方面所能提供的所有功能，这篇博文将帮助您深入了解如何利用这种强大且高度可定制的模型存储格式。

MLflow 提供了内置方法来记录和处理许多流行的机器学习和生成式 AI 框架以及模型提供商的模型，例如 scikit-learn、PyTorch、HuggingFace transformers 和 LangChain。例如，mlflow.sklearn.log_model 会将 scikit-learn 模型记录为 MLflow 工件，而无需您定义自定义的预测方法或处理工件的方法。

然而，在某些情况下，您可能正在使用 MLflow 没有内置方法的框架，或者您可能想要与模型的默认预测输出不同的内容。在这些情况下，MLflow 允许您创建自定义模型，以支持几乎任何框架，并将自定义逻辑集成到现有支持的框架中。

最简单的情况是，只需要定义一个自定义的 predict 方法并记录模型。以下示例定义了一个简单的 pyfunc 模型，该模型仅返回其输入的平方。

import mlflow

# Define a custom model
class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        # Directly return the square of the input
        return model_input**2


# Save the model
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=MyModel()
    )

# Load the model
loaded_model = mlflow.pyfunc.load_model(
    model_uri=model_info.model_uri
)

# Predict
loaded_model.predict(2)

让我们深入探讨这是如何工作的，首先从一些基本概念开始。

模型和模型风味 (Models and Model Flavors)

MLflow 模型是一个目录，其中包含在不同环境中重现机器学习模型所需的一切。除了存储的模型本身，最重要的组件是 `MLmodel` YAML 文件，该文件指定了模型支持的模型风味。模型风味 (model flavor) 是一组规则，用于指定 MLflow 如何与模型进行交互（即，保存、加载和从中获取预测）。

当您在 MLflow 中创建自定义模型时，它具有 `python_function` 或 pyfunc 模型风味，这是一种 MLflow 中跨格式的“通用翻译器”。当您使用内置的模型风味在 MLflow 中保存模型时，例如使用 mlflow.sklearn.log_model，该模型除了其特定于框架的风味外，还具有 pyfunc 模型风味。同时拥有特定于框架和 pyfunc 模型风味，可以让您通过框架的本机 API（例如，`scikit-learn`）或通过 pyfunc 风味的框架无关的推理 API 来使用模型。

具有 pyfunc 风度的模型会作为 mlflow.pyfunc.PyfuncModel 类的实例加载，该类公开了一个标准化的 predict 方法。这使得通过单个函数调用实现直接推理，而无需考虑底层模型的实现细节。

定义自定义 MLflow Pyfunc 模型 (Defining Custom MLflow Pyfunc Models)

将任何支持的机器学习框架的模型保存为 MLflow 模型，都会创建一个 pyfunc 模型风味，该风味为管理和使用模型提供了框架无关的接口。但是，如果您正在使用没有 MLflow 集成的框架，或者您正在尝试从模型中获得一些自定义行为呢？自定义 pyfunc 模型允许您支持几乎任何框架并将自定义逻辑集成到其中。

要实现自定义 pyfunc 模型，请定义一个继承自 PythonModel 类的新 Python 类并实现必要的方法。最低限度，这需要实现一个自定义的 predict 方法。接下来，创建模型的一个实例并记录或保存模型。一旦您加载了保存或记录的模型，就可以使用它进行预测。

让我们通过几个示例来学习，每个示例都增加了一些复杂性，并强调了定义自定义 pyfunc 模型的不同方面。我们将涵盖实现 pyfunc 模型中自定义行为的四种主要技术：

1. 实现自定义 `predict` 方法
2. 实现自定义 `__init__` 方法
3. 实现自定义 `load_context` 方法
4. 实现用户定义的自定义方法

定义自定义 `predict` 方法 (Defining a custom `predict` method)

最低限度，pyfunc 模型应指定一个自定义的 predict 方法，该方法定义了在调用 `model.predict` 时会发生什么。这是一个自定义模型示例，它对模型输入应用简单的线性变换，将每个输入乘以二并加上三。

import pandas as pd
import mlflow
from mlflow.pyfunc import PythonModel


# Custom PythonModel class
class SimpleLinearModel(PythonModel):
    def predict(self, context, model_input):
        """
        Applies a simple linear transformation
        to the input data. For example, y = 2x + 3.
        """
        # Assuming model_input is a pandas DataFrame with one column
        return pd.DataFrame(2 * model_input + 3)


with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="linear_model",
        python_model=SimpleLinearModel(),
        input_example=pd.DataFrame([10, 20, 30]),
    )

请注意，您也可以（也应该）在保存/记录模型时包含签名和输入示例。如果您提供输入示例，签名将自动推断。模型签名提供了一种 MLflow 强制执行模型正确用法的方式。

一旦我们定义了模型路径并保存了模型的实例，我们就可以加载保存的模型并使用它来生成预测。

# Now the model can be loaded and used for predictions
loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
model_input = pd.DataFrame([1, 2, 3])  # Example input data
print(loaded_model.predict(model_input))  # Outputs transformed data

这将返回

:    0
: 0  5
: 1  7
: 2  9

请注意，如果只需要一个自定义 `predict` 方法，即，如果您的模型没有任何需要特殊处理的工件，您可以直接保存或记录 `predict` 方法，而无需将其包装在 Python 类中。

import mlflow
import pandas as pd


def predict(model_input):
    """
    Applies a simple linear transformation
    to the input data. For example, y = 2x + 3.
    """
    # Assuming model_input is a pandas DataFrame with one column
    return pd.DataFrame(2 * model_input + 3)


# Pass predict method as python_model argument to save/log model
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="simple_function",
        python_model=predict,
        input_example=pd.DataFrame([10, 20, 30]),
    )

请注意，使用此方法，我们**必须包含**输入示例以及自定义的 predict 方法。我们还必须修改 predict 方法，使其只接受一个输入（即，没有 self 或 context）。运行此示例，然后使用与前面代码块相同的代码加载，将保留与使用类定义的示例相同的输出。

参数化自定义模型 (Parameterizing the custom model)

现在，假设我们想要参数化自定义线性函数模型，以便可以使用不同的斜率和截距。我们可以定义 `__init__` 方法来设置自定义参数，如下例所示。请注意，自定义模型类的 `__init__` 方法不应用于加载外部资源，例如数据文件或预训练模型；这些将在 `load_context` 方法中处理，我们稍后会讨论。

import pandas as pd
import mlflow
from mlflow.pyfunc import PythonModel


# Custom PythonModel class
class ParameterizedLinearModel(PythonModel):
    def __init__(self, slope, intercept):
        """
        Initialize the parameters of the model. Note that we are not loading
        any external resources here, just setting up the parameters.
        """
        self.slope = slope
        self.intercept = intercept

    def predict(self, context, model_input):
        """
        Applies a simple linear transformation
        to the input data. For example, y = 2x + 3.
        """
        # Assuming model_input is a pandas DataFrame with one column
        return pd.DataFrame(self.slope * model_input + self.intercept)


linear_model = ParameterizedLinearModel(10, 20)

# Saving the model with mlflow
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="parameter_model",
        python_model=linear_model,
        input_example=pd.DataFrame([10, 20, 30]),
    )

同样，我们可以加载此模型并进行一些预测。

loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
model_input = pd.DataFrame([1, 2, 3])  # Example input data
print(loaded_model.predict(model_input))  # Outputs transformed data

:     0
: 0  30
: 1  40
: 2  50

在许多情况下，我们可能希望以这种方式参数化模型。我们可以在 `__init__` 方法中定义变量来：

• 设置模型超参数。
• 使用不同参数集进行模型的 A/B 测试。
• 设置用户特定的自定义项。
• 切换功能。
• 设置，例如，访问外部 API 的模型的访问凭据和端点。

在某些情况下，我们可能希望在推理时而不是在初始化模型时传递参数。这可以通过模型推理参数来实现。要使用推理参数，我们必须传递包含 `params` 的有效模型签名。下面介绍如何将前面的示例改编为使用推理参数。

import pandas as pd
import mlflow
from mlflow.models import infer_signature
from mlflow.pyfunc import PythonModel


# Custom PythonModel class
class LinearFunctionInferenceParams(PythonModel):
    def predict(self, context, model_input, params):
        """
        Applies a simple linear transformation
        to the input data. For example, y = 2x + 3.
        """
        slope = params["slope"]
        # Assuming model_input is a pandas DataFrame with one column
        return pd.DataFrame(slope * model_input + params["intercept"])


# Set default params
params = {"slope": 2, "intercept": 3}

# Define model signature
signature = infer_signature(model_input=pd.DataFrame([10, 20, 30]), params=params)

# Save the model with mlflow
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="model_with_params",
        python_model=LinearFunctionInferenceParams(),
        signature=signature,
    )

像往常一样加载模型后，您现在可以将 `params` 参数传递给 `predict` 方法，从而使您可以使用相同的加载模型来进行不同的参数组合。

loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)

parameterized_predictions = loaded_model.predict(
    pd.DataFrame([10, 20, 30]), params={"slope": 2, "intercept": 10}
)
print(parameterized_predictions)

:     0
: 0  30
: 1  50
: 2  70

使用 `load_context` 加载外部资源 (Loading external resources with `load_context`)

自定义模型通常需要外部文件，例如模型权重，以便执行推理。这些文件或工件必须小心处理，以避免不必要地将文件加载到内存或在模型序列化过程中出现错误。在构建 MLflow 中的自定义 pyfunc 模型时，您可以使用 `load_context` 方法来正确处理模型工件。

`load_context` 方法接收一个 `context` 对象，其中包含模型在推理期间可以使用的工件。您可以使用保存或记录模型时的 `artifacts` 参数来指定这些工件，从而使它们可以通过 `context.artifacts` 字典在 `load_context` 方法中访问。

在实践中，`load_context` 方法通常通过处理模型工件的加载来初始化 `predict` 方法调用的模型。

这引出了一个重要问题：为什么我们在 `load_context` 方法中加载工件并定义模型，而不是在 `__init__` 或直接在 `predict` 中？正确使用 `load_context` 对 MLflow pyfunc 模型的可维护性、效率、可扩展性和可移植性至关重要。这是因为：

• `load_context` 方法在通过 `mlflow.pyfunc.load_model` 加载模型时执行一次。此设置确保在此方法中定义的资源密集型过程（例如加载大型模型文件）不会不必要地重复。如果工件加载在 predict 方法中完成，它将在每次进行预测时发生。这对于资源密集型模型来说效率极低。
• 保存或记录 MLflow `pyfunc` 模型涉及序列化 Python 模型类（您创建的 `mlflow.pyfunc.PythonModel` 的子类）及其属性。复杂的 ML 模型并非总是与用于序列化 Python 对象的方法兼容，如果它们作为 Python 对象的属性创建，可能会导致错误。

例如，假设我们想加载 `gguf` 模型格式（一种用于存储模型以进行推理的文件格式）的大型语言模型 (LLM)，并使用 ctransformers 库运行它。截至撰写本文时，没有内置的模型风味允许我们使用 `gguf` 模型进行推理，因此我们将创建一个自定义 pyfunc 模型，该模型在 `load_context` 方法中加载所需的库和模型文件。具体来说，我们将加载 Mistral 7B 模型的 AWQ 版本的量化版本。

首先，我们将使用 huggingface hub cli 下载模型快照。

huggingface-cli download TheBloke/Mistral-7B-v0.1-GGUF \
                mistral-7b-v0.1.Q4_K_M.gguf \
                --local-dir /path/to/mistralfiles/ \
                --local-dir-use-symlinks False

然后我们将定义我们的自定义 `pyfunc` 模型。请注意添加了 `load_context` 方法。

import ctransformers
from mlflow.pyfunc import PythonModel


class CTransformersModel(PythonModel):
    def __init__(self, gpu_layers):
        """
        Initialize with GPU layer configuration.
        """
        self.gpu_layers = gpu_layers
        self.model = None

    def load_context(self, context):
        """
        Load the model from the specified artifacts directory.
        """
        model_file_path = context.artifacts["model_file"]

        # Load the model
        self.model = ctransformers.AutoModelForCausalLM.from_pretrained(
            model_path_or_repo_id=model_file_path,
            gpu_layers=self.gpu_layers,
        )

    def predict(self, context, model_input):
        """
        Perform prediction using the loaded model.
        """
        if self.model is None:
            raise ValueError(
                "The model has not been loaded. "
                "Ensure that 'load_context' is properly executed."
            )
        return self.model(model_input)

这里有很多内容，让我们分解一下。以下是关键点：

• 和以前一样，我们使用 `__init__` 方法来参数化模型（在这种情况下，是为了设置模型的 `gpu_layers` 参数）。
• `load_context` 方法的目的是加载 `predict` 方法使用所需的工件。在这种情况下，我们需要加载模型及其权重。
• 您会注意到我们引用了 `context.artifacts["model_file"]`。这来自于 `mlflow.pyfunc.save_model` 或 `mlflow.pyfunc.load_model` 的 `artifacts` 参数，如下面的代码片段所示。这是处理 `pyfunc` 模型的重要组成部分。`predict` 和 `load_context` 方法可以通过 `context.artifacts` 对象访问 `save_model` 或 `log_model` 方法的 `artifacts` 参数中定义的工件。`load_context` 在通过 `load_model` 加载模型时执行；如前所述，这提供了一种确保潜在耗时的模型初始化不会在每次使用模型进行预测时都发生的方法。

现在我们可以初始化并保存模型的一个实例。请注意 `save_model` 函数的 `artifacts` 参数。

# Create an instance of the model
mistral_model = CTransformersModel(gpu_layers=50)

# Log the model using mlflow with the model file as an artifact
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        artifact_path="mistral_model",
        python_model=mistral_model,
        artifacts={"model_file": model_file_path},
        pip_requirements=[
            "ctransformers==0.2.27",
        ],
    )

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

# Make a prediction with the model
loaded_model.predict("Question: What is the MLflow Pyfunc model flavor?")

总结：正确使用 `load_context` 方法有助于确保模型工件的高效处理，并防止由于尝试将工件定义为模型属性而可能导致的序列化错误。

定义自定义方法 (Defining custom methods)

您可以在自定义 `pyfunc` 模型中定义自己的方法来处理诸如预处理输入或后处理输出之类的任务。然后，这些自定义方法可以被 predict 方法调用。请记住，与 `__init__` 和 `predict` 一样，这些自定义方法**不应用于**加载工件。加载工件是 `load_context` 方法的专属职责。

例如，我们可以修改 `CTransformersModel` 来加入一些提示格式设置，如下所示。

import ctransformers
from mlflow.pyfunc import PythonModel


class CTransformersModel(PythonModel):
    def __init__(self, gpu_layers):
        """
        Initialize with GPU layer configuration.
        """
        self.gpu_layers = gpu_layers
        self.model = None

    def load_context(self, context):
        """
        Load the model from the specified artifacts directory.
        """
        model_file_path = context.artifacts["model_file"]

        # Load the model
        self.model = ctransformers.AutoModelForCausalLM.from_pretrained(
            model_path_or_repo_id=model_file_path,
            gpu_layers=self.gpu_layers,
        )

    @staticmethod
    def _format_prompt(prompt):
        """
        Formats the user's prompt
        """
        formatted_prompt = (
            "Question: What is an MLflow Model?\n\n"
            "Answer: An MLflow Model is a directory that includes "
            "everything that is needed to reproduce a machine "
            "learning model across different environments. "
            "It is essentially a container holding the trained model "
            "files, dependencies, environment details, input examples, "
            "and additional metadata. The directory also includes an "
            "MLmodel YAML file, which describes the different "
            f"flavors of the model.\n\nQuestion: {prompt}\nAnswer: "
        )

        return formatted_prompt

    def predict(self, context, model_input):
        """
        Perform prediction using the loaded model.
        """
        if self.model is None:
            raise ValueError(
                "Model was not loaded. Ensure that 'load_context' "
                "is properly executed."
            )
        return self.model(self._format_prompt(model_input))

现在 `predict` 方法可以访问私有的 `_format_prompt` 静态方法来应用自定义的提示格式。

依赖项和源代码 (Dependencies and Source Code)

上面定义的自定义 `CTransformersModel` 使用 `ctransformers` 库。有几种方法可以确保此库（以及任何其他源代码，包括本地设备上的源代码）与模型正确加载。正确指定依赖项对于确保自定义模型在不同环境中按预期工作至关重要。

有三种主要的方法需要注意，用于指定依赖项：

• 使用 `save_model` 或 `log_model` 的 `pip_requirements` 参数显式定义 pip 要求。
• 使用 `save_model` 或 `log_model` 的 `extra_pip_requirements` 参数向自动生成的要求集添加额外的 pip 要求。
• 使用 `save_model` 或 `log_model` 的 `conda_env` 参数定义 Conda 环境。

之前，我们使用了第一种方法来指定需要 `ctransformers` 库。

# Log the model using mlflow with the model file as an artifact
with mlflow.start_run():
    model_info = mlflow.pyfunc.save_model(
        artifact_path="mistralmodel",
        python_model=mistral_model,
        artifacts={"model_file": "path/to/mistral/model/on/local/filesystem"},
        pip_requirements=[
            "ctransformers==0.2.27",
        ],
    )

如果您不显式指定依赖项，MLflow 将尝试推断正确的要求集和环境详细信息。为了提高准确性，**强烈建议**在保存或记录模型时包含 `input_example`，因为内部会执行样本推理步骤，该步骤将捕获与推理执行相关的任何已加载库引用，从而提高记录正确依赖项的可能性。

您还可以使用 `code_path` 参数处理本地文件系统上的自定义代码。`code_path` 接受指向 Python 文件依赖项的路径列表，并在模型加载之前将其添加到系统路径中，以便自定义 pyfunc 模型可以从这些模块导入。

有关 `pip`、`Conda` 和本地代码要求的接受格式的更多详细信息，请参阅 log_model 和 save_model 函数的文档。

总结：MLflow 中的自定义 Pyfunc 模型 (Summary: Custom Pyfunc Models in MLflow)

MLflow 提供了内置方法来处理许多流行的机器学习框架的模型，例如scikit-learn、PyTorch和Transformers。当您想处理尚无内置模型风向的模型，或者想为具有内置模型风向的模型实现自定义 predict 方法时，可以定义自己的自定义 `mlflow.pyfunc` 模型。

有几种方法可以自定义 `pyfunc` 模型以获得所需的行为。最低限度，您可以实现自定义 `predict` 方法。如果您的模型需要保存或加载工件，您还应该实现 `load_context` 方法。为了进一步定制，您可以使用 `__init__` 方法来设置自定义属性，并定义自己的自定义方法来进行预处理和后处理。结合这些方法，您可以灵活地为机器学习模型定义自定义逻辑。

深入学习 (Further Learning)

有兴趣了解更多关于自定义 `pyfunc` 实现的信息吗？您可以访问：

• 使用 MLflow 为高级 LLM 构建自定义 Pyfunc
• 为传统 ML 构建自定义 Python 函数模型
• 自定义 PyFunc 笔记本示例