跳到主要内容

理解 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 函数模型 (Python Function Flavor)：这是 MLflow Python 模型的默认模型接口。它确保可以使用一致的 API 加载和交互所有 MLflow Python 模型。

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

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

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

自定义 Pyfunc 模型的强大之处​

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

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

最适合自定义 Pyfunc 的情况​

在许多情况下，自定义 Pyfunc 都非常宝贵：

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

    • 在 Apache Spark 或 Ray 等分布式系统中，推理是跨多个核心并行化的，存在为每个核心加载多个模型副本的风险。这会给系统的资源带来巨大压力，尤其是对于大型模型。
    • 通过自定义 Pyfunc，您可以确保每个工作节点或执行器只加载模型的一个副本，从而优化资源利用率并加快推理速度。

  2. 不支持的模型:

    • 虽然 MLflow 为流行框架提供了广泛的命名模型，但机器学习生态系统非常庞大。可能存在一些尚不支持的小众或新兴框架。
    • 自定义 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. 构件映射:

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

  4. Python 模型初始化:

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

  5. 上下文加载:

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

  6. 模型准备就绪:

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

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

后续步骤​

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

探索教程笔记本