理解 MLflow 中的 PyFunc

在 MLflow 领域，虽然命名的风格提供了针对流行框架的特定功能，但有些情况和需求超出了这些预定义的路径。这时，自定义的 :py:mod:`pyfunc <mlflow.pyfunc>`（Python 函数）作为一个通用接口，使你能够通过定义一个自定义的 Python 函数，将任何框架的模型封装到 MLflow 模型中。

PyFunc 版本的模型与其他任何 MLflow 模型类型的交互方式相同，提供了 save_model() 和 log_model() 接口，分别用于创建（保存）和访问（加载）模型。

因为这些自定义模型包含 python_function 风格，它们可以部署到任何 MLflow 支持的生产环境中，例如 SageMaker、AzureML、Databricks、Kubernetes 或本地 REST 端点。

为什么选择 PyFunc？

  1. 灵活性: 它提供了与任何机器学习库或框架一起工作的自由，确保了MLflow对广泛用例的适应性。

  2. 统一接口：通过 pyfunc，您可以获得一致的 API。一旦您的模型符合此接口，您就可以利用 MLflow 的所有部署工具，而无需担心底层框架。

  3. 自定义逻辑：除了模型本身，pyfunc 还允许进行预处理和后处理，增强了模型的部署能力。

PyFunc 的组成部分

  1. Python 函数风格: 这是 MLflow Python 模型的默认模型接口。它确保每个 MLflow Python 模型都可以使用一致的 API 加载和交互。

  2. 文件系统格式: 一个包含所有必需数据、代码和配置的结构化目录，确保封装的模型及其依赖项是自包含且可复现的。

  3. MLModel 配置: 一个重要的描述符，MLmodel 文件提供了关于模型的详细信息，包括其加载模块、代码、数据和环境。

  4. 自定义Pyfunc模型：一个强大的功能，超越了命名风格，允许创建具有自定义逻辑、数据转换等的模型。

自定义 Pyfunc 模型的力量

虽然 MLflow 的命名风格为许多框架提供了开箱即用的解决方案，但它们可能无法满足所有需求。这就是自定义 pyfunc 模型的亮点所在。无论你是使用小众框架，需要实现专门的预处理，还是希望集成推理后的逻辑，自定义 pyfunc 模型都提供了相应的工具来实现这些功能。

通过定义一个继承自 PythonModel 的 Python 类并实现必要的方法，你可以创建一个定制的 pyfunc 模型，以满足你的特定需求。

自定义 Pyfunc 可能最适用的条件

在许多场景中，自定义的 Pyfunc 变得不可或缺：

  1. 使用大型模型进行分布式推理：

    • 在像Apache Spark或Ray这样的分布式系统中，推理在多个核心上并行化，每个核心加载一个模型的副本存在风险。这会显著消耗系统资源，尤其是在使用大型模型时。

    • 通过自定义 Pyfunc，您可以确保每个工作节点或执行器仅加载模型的一个副本，从而优化资源使用并加快推理速度。

  2. 不支持的模型：

    • 虽然 MLflow 为流行的框架提供了多种命名风格，但机器学习生态系统非常庞大。可能存在尚未支持的小众或新兴框架。

    • 自定义Pyfunc提供了一种无缝封装和管理来自任何不支持框架的模型的方法。

  3. 自定义推理方法：

    • 默认的 .predict() 方法可能并不总是满足特定的需求。也许你需要一个能生成对数、不确定性或其他指标的方法。

    • 自定义的 Pyfunc 可以封装任何推理方法，确保部署的模型行为完全符合需求。

  4. 加载辅助数据或外部系统：

    • 有时，模型的推理不仅仅关乎模型本身。它可能需要引用未与模型一起保存的外部数据，或者可能需要连接到其他系统。

    • 考虑一个场景，模型在预测时需要查询向量数据库中的条目。自定义的 Pyfunc 可以利用 load_context 方法来加载配置文件。这为自定义的 predict 方法提供了配置数据，使其能够在推理期间连接到外部服务。

自定义 Pyfunc 的内部工作原理

理解在调用 mlflow.pyfunc.load_model() 期间的事件序列对于充分利用自定义 Pyfuncs 至关重要。以下是加载自定义 pyfunc 时发生的事件序列的逐步分解，以及在保存模型时如何声明覆盖，以及这些覆盖如何被访问和引用以控制加载模型对象的行为。

标签、实验和运行关系

Pyfunc 加载过程

  1. 启动：

    • 当调用 mlflow.pyfunc.load_model() 时，过程开始，表明意图加载一个自定义的 Pyfunc 模型以供使用。

  2. 模型配置检索:

    • 系统获取与保存的模型关联的 MLmodel 配置文件。该描述符提供了关于模型的基本细节，包括其加载模块、代码、数据和环境。

  3. 工件映射：

    • 保存的模型工件，可能包括序列化的模型对象、辅助数据或其他必要文件，会被映射。这种映射确保自定义的 Pyfunc 知道在哪里找到它所需的一切。

  4. Python 模型初始化:

    • 定义自定义 Pyfunc 的 Python 类（通常继承自 PythonModel）被初始化。在这个阶段，模型尚未准备好进行推理，但已准备好进行后续的加载步骤。

  5. 上下文加载：

    • 自定义 Pyfunc 的 load_context 方法被调用。此方法旨在加载任何外部引用或执行初始化任务。例如，它可以反序列化一个模型对象，加载用于连接外部服务的配置文件，或准备模型所需的其他任何资源。

  6. 模型就绪:

    • 加载上下文后，自定义的 Pyfunc 模型现已完全初始化，并准备好进行推理。对其 predict 方法的任何后续调用现在都将执行在其内部定义的自定义逻辑，从而产生设计的结果。

值得注意的是，这个序列确保了自定义的 Pyfunc 模型一旦加载，就是一个完全自包含的单元，不仅封装了模型，还包括了任何自定义逻辑、数据转换以及它所需的外部引用。这种设计确保了无论模型部署在哪里，都能保证可重复性和一致性。

下一步

既然你已经理解了 pyfunc 的重要性和组成部分，下一步就是深入了解如何构建它们。

Explore the tutorial notebooks