跳到主要内容

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,您可以在自定义类中嵌入任何库的任何代码或模型,同时享受与命名风格相关的统一性优势。

要更深入地了解这些功能,我们来研究 MLflow 模型 的核心结构。

MLflow 中模型的组件

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

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

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

Components of a Model in MLflow

MLflow 中模型的基本组件

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

  2. 附加二进制文件:对于某些模型,可能需要辅助文件。例如,NLP 模型的 tokenizer、预处理的 scaler,甚至是非参数元素,如决策树或 k-means 聚类中心。

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

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

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

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

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

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

Components of a Model in the MLflow UI

MLflow Artifact 查看器中显示的模型组件

注意

MLflow 中日志模型目录的内容取决于保存或记录模型时提供的可选参数数量以及底层基础模型类型。某些模型风格比其他风格具有额外的元数据和序列化工件。

在创建自定义 PyFunc 时,理解此处显示的组件非常重要,因为在创建和使用自定义 PyFunc 时,您将与这种结构及其中的元素进行交互。

理解“命名风格”

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

命名风格的主要特性包括:

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

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

  • 自动日志记录:某些命名风格支持自动日志记录,这是一项在训练过程完成后自动记录模型工件和训练元数据的功能。

命名风格的特点

命名风格封装了多项功能:

  • 统一 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 加载的集成。