跳到主要内容

MLflow 中的模型、风格和 PyFunc

在 MLflow 生态系统中，“flavors”（风味）在模型管理中起着至关重要的作用。本质上，“flavor”是特定机器学习库的指定封装器。例如，即使 spark-ml 包会生成各种模型类型，如 Pipeline、LogisticRegressionModel 或 RandomForestModel，它们都归属于 Spark flavor 的伞下。这种抽象确保了，无论模型的底层结构如何，其 Spark flavor 的变体都可以使用 MLflow 的命名 flavor 工具无缝地保存、记录和检索。

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

MLflow 的 flavor 设计确保了一定程度的统一性。对于每个库，其对应的 MLflow flavor 定义了用于推理部署的已加载 pyfunc 的行为。每个 flavor 都规定了 predict 方法的行为，确保了一种一致但又略显僵化的格式。

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

MLflow's sklearn implementation

MLflow sklearn flavor 所采用标准的示意图

尽管 MLflow 努力为每种 flavor 提供一个普遍适用的 pyfunc 表示，但并非总是能够适应特定库生成的每种独特的模型场景。

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

为了深入了解这些功能，让我们来检查 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 的 artifact 查看器中查看已保存的模型时，可以看到所有这些元素。

Components of a Model in the MLflow UI

在 MLflow artifact 查看器中看到的模型组件

注意

MLflow 中记录的模型目录的内容取决于在保存或记录模型时提供的可选参数数量，以及底层基础模型类型。一些模型 flavor 相较于其他 flavor 拥有额外的元数据和序列化 artifact。

此处显示的组件对于创建自定义 pyfunc 至关重要，因为您在创建和使用自定义 PyFuncs 时将与此结构及其内部元素进行交互。

理解“命名 Flavors”​

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

命名 flavor 的关键属性包括：

  • 根命名空间集成：命名 flavor 可以直接从 MLflow 根命名空间访问，从而实现直接交互。

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

  • 自动记录：某些命名 flavor 支持自动记录，这是一项功能，可在训练过程完成后自动记录模型 artifact 和训练元数据。

命名 Flavors 的特点​

命名 flavor 封装了多项功能：

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

  • 维护与可靠性：作为 MLflow 项目的一部分，命名 flavor 经过核心维护者的严格测试和更新。

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

  • 自定义 Python 函数包装器：每种 flavor 都包含一个特定的实现，将底层框架的方法映射到标准的 Python 函数，并对函数的行为做出某些决策。

  • 简化的高级 API：尽管有能力处理复杂的细节，但命名 flavor 的高级 API 设计得易于使用。

包含为命名 Flavor 的标准​

将一个框架作为命名 flavor 包含在 MLflow 中并非随意。

标准包括：

  • 受欢迎程度和需求：拥有大量行业采用的框架更受青睐。包含的决定还取决于用户请求的频率以及更广泛的 ML 社区中的感知需求。

  • 框架稳定性：命名 flavor 通常与稳定的、有活跃维护且没有过于复杂或限制性构建要求的框架相关联，这些要求可能会导致集成它们成为一项不可能的任务。

命名 Flavors 的结构​

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

  • get_default_conda_env()：返回 flavor 所需的 conda 依赖项列表。

  • get_default_pip_requirements()：列出 flavor 所必需的 PyPI 依赖项。

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

  • save_model()：管理序列化过程，确保模型、其元数据和其他相关 artifact 被适当地存储。

  • log_model()：这是 save_model() 的扩展版本，除了保存过程外，还便于模型注册。

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

处理 MLflow 中不受支持的模型​

对于未被支持为命名 flavor 的机器学习框架，MLflow 提供了定义自定义 PyFuncs 的灵活性。

本教程将引导您完成此过程，使您能够将几乎任何模型整合到 MLflow 生态系统中。

创建可重用的自定义 Flavors​

对于那些在不同项目中频繁使用特定自定义 PyFuncs 的用户，MLflow 的架构支持通过插件式接口开发自定义 flavors。虽然本教程的范围不包括此主题的详尽指南，但一般的方法包括创建一个包含模型类型保存、加载和记录功能的模块。然后，将创建一个 PyFunc 包装器类，为将自定义 flavor 加载为 PyFunc 提供集成。