跳到主要内容

理解 MLflow 中的 PyFunc

在 MLflow 领域，虽然命名风格 (named flavors) 为流行的框架提供了特定的功能，但有些情况和需求超出了这些预设的路径。此时就需要自定义的 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() 方法可能无法始终满足特定需求。您可能需要一个能生成 logit、不确定性或其他指标的方法。
    • 自定义 Pyfunc 可以封装任何推理方法，确保部署的模型完全按照需要运行。

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

    • 有时，模型的推理不仅仅依赖于模型本身。它可能需要引用未随模型一起保存的外部数据，或者可能需要连接到其他系统。
    • 考虑这样一个场景：模型在预测过程中需要在向量数据库中查找条目。自定义 Pyfunc 可以利用 load_context 方法加载配置文件。这为自定义 predict 方法提供了配置数据，使其能够在推理过程中连接到外部服务。

自定义 PyFunc 的内部工作原理​

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

Tags, experiments, and runs relationships

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 的重要性和组成部分，下一步就是深入了解如何构建它们。

探索教程 Notebook