跳到主要内容

MLflow 中的模型、风格和 PyFunc

在 MLflow 生态系统中,“模型风味”(flavors)在模型管理中扮演着关键角色。本质上,“模型风味”是针对特定机器学习库的指定封装器。例如,尽管 spark-ml 包会生成多种模型类型,例如 PipelineLogisticRegressionModelRandomForestModel,但它们都归属于 Spark 模型风味。这种抽象确保了,无论模型的底层结构如何,其 Spark 模型风味变体都可以使用 MLflow 的命名模型风味工具无缝地进行保存、记录和检索。

模型风味简化了在不同框架中保存、加载和处理机器学习模型的过程。它们考虑了每个库在模型序列化和反序列化方面的独特方法。

MLflow 的模型风味设计确保了一定程度的统一性。对于每种库,其相应的 MLflow 模型风味定义了加载后的 pyfunc 在进行推理部署时的行为。每种模型风味都规定了 predict 方法的行为,从而确保了一种一致但略显僵化的格式。

为了理解这些限制,可以 sklearn 模型风味为例。下图描绘了它的实现,重点介绍了 MLflow 标准化的 API 和序列化方法。

MLflow's sklearn implementation

MLflow sklearn 模型风味所采纳标准的示意图

虽然 MLflow 致力于为每种模型风味提供一个通用的 pyfunc 表示,但并非总是能够适应特定库生成的每一种独特的模型场景。

然而,这也有一个好消息。MLflow 提供了通过扩展 PythonModel 基础类来创建自定义 pyfunc 的灵活性,该类是所有命名模型风味的 pyfunc 变体的基础。通过正确实现 PythonModel,您可以将任何库中的任何代码或模型嵌入自定义类中,同时享受与命名模型风味相关的统一性优势。

为了更深入地了解这些功能,让我们 examine MLflow Model 的核心结构。

MLflow 中模型组件

当提到“模型”时,大多数从业者会想到机器学习训练过程中学习到的参数或权重。这些通常被保存为一个文件或一组文件,然后用于对新、未见过的数据进行预测。然而,在 MLOps 领域,尤其是在 MLflow 中,“模型”的概念更为广泛。

在 MLflow 中,模型不仅仅是包含学习参数的二进制文件。它是一个完整的包或捆绑,封装了在各种环境中可靠地重现预测所需的一切。

这包括模型的权重,但远不止于此。

Components of a Model in MLflow

MLflow 中模型的基本组成部分

  1. 模型二进制文件:这是核心部分——实际保存的模型权重或参数。这是许多人认为的“模型”。

  2. 额外的二进制文件:某些模型可能需要辅助文件。例如,NLP 模型的分词器、预处理的缩放器,甚至非参数化元素,如决策树或 K-Means 质心。

  3. 预加载代码:某些模型可能需要在推理环境中加载自定义代码。这可能是用于预处理、后处理或其他自定义逻辑。

  4. 库依赖项:为了使模型能够正常运行,它可能依赖于特定版本的库。MLflow 会跟踪这些依赖项,确保模型运行的环境与训练它的环境相匹配。

  5. 元数据:这包含有关模型 lineage 的重要信息。它可以跟踪模型由谁训练、使用什么代码、何时何地等细节。这些元数据对于模型治理、审计和可重现性至关重要。

  6. PyFunc 签名:为了确保无缝部署和推理,MLflow 将模型封装在标准化的 pyfunc 接口中。此接口定义了预期的输入和输出格式,确保一致性。

  7. 输入示例:这是一个可选组件,提供了一个可用于测试的示例输入,以确保已部署的模型正常工作。

当在 MLflow UI 中查看已保存的模型时,所有这些元素都可以在制品查看器中看到。

Components of a Model in the MLflow UI

在 MLflow 制品查看器中看到模型组件

注意

MLflow 中记录的模型目录的内容取决于保存或记录模型时提供的可选参数的数量以及底层基本模型类型。与某些模型相比,一些模型风味具有额外的元数据和序列化制品。

这里显示的组件对于创建自定义 pyfunc 非常重要,因为当你创建和使用自定义 PyFuncs 时,这个结构及其内部元素就是你将要交互的对象。

理解“命名风味” (Named Flavors)

MLflow 中的命名风味 (Named Flavor) 指的是与特定机器学习或数据处理框架关联的预定义实体。例如,如果你正在使用 Scikit-Learn 模型,你可能会使用 mlflow.sklearn.save_model()mlflow.sklearn.load_model()mlflow.sklearn.log_model() 等方法。

命名风味的关键属性包括

  • 根命名空间集成:命名风味可以直接从 MLflow 根命名空间访问,从而实现直接交互。

  • PyFunc 兼容性:使用命名风味保存的模型可以作为 PyFunc 加载回来。这有助于与各种部署环境集成,无论是实时推理平台、基于 Spark 的批处理,还是任何能够调用 Python 函数的系统。

  • 自动记录 (Autologging):某些命名风味支持自动记录功能,该功能可在训练过程完成后自动记录模型制品和训练元数据。

命名风味的特点

命名风味封装了多项功能

  • 统一 API:尽管底层机器学习框架存在差异,但命名风味为模型保存、加载和记录提供了一套一致的方法。这种一致性延伸到高级功能,如签名声明、输入示例存储、自定义依赖项和模型注册。

  • 维护与可靠性:作为 MLflow 项目的一部分,命名风味会经过核心维护者的严格测试和更新。

  • 序列化方法:每种命名风味都利用与其关联框架相关的原生序列化机制。

  • 自定义 Python 函数封装器:每种风味都包含一个特定的实现,该实现将底层框架的方法映射到标准的 Python 函数,并就函数的行为做出某些决策。

  • 简化的高级 API:尽管它们能够处理复杂的细节,但命名风味的高级 API 设计简单易用。

作为命名风味包含的标准

将一个框架包含在 MLflow 中作为命名风味不是任意的。

标准包括

  • 流行度和需求:行业采用率较高的框架更受欢迎。包含还取决于用户请求的频率以及更广泛的 ML 社区中的预期需求。

  • 框架稳定性:命名风味通常与稳定、有积极维护且没有过于复杂或限制性构建要求的框架相关联,这些要求可能会使集成变得不可能。

命名风味的结构

MLflow 中的每个命名风味通常都实现了一组核心函数

  • get_default_conda_env():返回风味所需的 conda 依赖项列表。

  • get_default_pip_requirements():列出风味至关重要的 PyPI 依赖项。

  • load_model():处理反序列化过程,通过提供的可解析的 model_uri 从给定的制品存储中检索模型实例。

  • save_model():管理序列化过程,确保模型、其元数据和其他关联的制品得到妥善存储。

  • log_model()save_model() 的扩展版本,除了保存过程外,还促进模型注册。

此外,为了确保风味模型可以作为通用 Python 函数加载,需要一个 Wrapper 类来与 mlflow.pyfunc.load_model() 集成。

处理 MLflow 中不受支持的模型

对于未作为命名风味支持的机器学习框架,MLflow 提供了定义自定义 PyFuncs 的灵活性。

本教程将指导您完成此过程,使您能够将几乎任何模型集成到 MLflow 生态系统中。

创建可重用的自定义风味

对于那些经常在不同项目中使用特定自定义 PyFuncs 的用户,MLflow 的架构支持通过插件式接口开发自定义风味。虽然关于此主题的全面指南超出了本教程的范围,但通用方法包括创建一个包含保存、加载和记录模型类型的函数的模块。然后创建一个 PyFunc 包装器类,为将自定义风味加载为 PyFunc 提供集成。