跳到主要内容

理解 MLflow 中的 PyFunc

在 MLflow 的领域中，虽然命名 Flavor 提供了针对流行框架的特定功能，但有时也会遇到预定义路径之外的情况和需求。这时就需要自定义 pyfunc（Python 函数），它是一个通用接口，使您能够通过定义自定义 Python 函数，将任何框架的模型封装到 MLflow 模型中。

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

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

为什么选择 PyFunc？​

  1. 灵活性：它提供了与任何机器学习库或框架进行集成的自由，确保 MLflow 能够适应广泛的用例。

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

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

PyFunc 的组成部分​

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

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

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

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

自定义 Pyfunc 模型的力量​

虽然 MLflow 的命名 Flavor 为许多框架提供了开箱即用的解决方案，但它们可能无法满足所有需求。这时自定义 pyfunc 模型就派上用场了。无论您是使用一个小众框架，需要实现专门的预处理，还是想集成推理后的逻辑，自定义 pyfunc 模型都提供了实现这些功能的工具。

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

何时自定义 Pyfunc 可能最佳​

在许多情况下，自定义 Pyfunc 都非常有价值

  1. 具有大型模型的分布式推理:

    • 在 Apache Spark 或 Ray 等分布式系统中，推理会跨多个核心并行化，存在加载模型多个副本的风险，每个核心一个副本。这可能会严重占用系统资源，尤其是在处理大型模型时。
    • 通过自定义 Pyfunc，您可以确保每个工作节点或执行器只加载模型的一个副本，从而优化资源使用并加快推理速度。

  2. 不支持的模型:

    • 虽然 MLflow 为流行框架提供了广泛的命名 Flavor，但机器学习生态系统是庞大的。可能存在一些尚未得到支持的小众或新兴框架。
    • 自定义 Pyfunc 提供了一种无缝封装和管理任何此类不受支持框架的模型的方式。

  3. 自定义推理方法:

    • 默认的 .predict() 方法并不总是能满足特定需求。也许您需要一个能生成 logits、不确定性或其他指标的方法。
    • 自定义 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. Artifact 映射:

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

  4. Python 模型初始化:

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

  5. Context 加载:

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

  6. 模型就绪:

    • 在加载了 Context 之后，自定义 Pyfunc 模型现已完全初始化并准备好进行推理。之后对其 predict 方法的任何调用都将执行其中定义的自定义逻辑，生成预期的结果。

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

后续步骤​

现在您已经了解了 pyfunc 的重要性和组成部分，下一步是深入了解如何构建它们。

探索教程 Notebooks