跳到主要内容

MLflow 模型注册表

MLflow 模型注册表是一个集中的模型存储库、一组 API 和一个 UI，旨在协同管理模型的整个生命周期。它提供血缘（即哪个 MLflow 实验和运行生成了模型）、版本控制、别名、元数据标记和注释支持，以确保您在从开发到生产部署的每个阶段都拥有全面的信息。

为什么选择模型注册表？​

随着机器学习项目的复杂性和规模的增长，在不同环境、团队和迭代中手动管理模型变得越来越容易出错且效率低下。MLflow 模型注册表通过提供一个集中式、结构化的系统来组织和管理 ML 模型在其整个生命周期中的应用，从而解决了这一挑战。

使用模型注册表具有以下优点：

  • 🗂️ 版本控制：注册表自动跟踪每个模型的版本，允许团队比较迭代、回滚到以前的状态，并并行管理多个版本（例如，预生产与生产）。
  • 🧬 模型血缘和可追溯性：每个注册的模型版本都链接到生成它的 MLflow 运行、日志模型或笔记本，从而实现完全可重现性。您可以准确追溯模型的训练方式、使用的数据和参数。
  • 🚀 生产就绪工作流：模型别名（例如 @champion）和标签等功能使部署工作流更易于管理，以受控且可审计的方式将模型提升到实验、预生产或生产环境。
  • 🛡️ 治理和合规性：通过结构化元数据、标记和基于角色的访问控制（与 Databricks 或托管 MLflow 服务等后端一起使用时），模型注册表支持企业级 ML 操作所需的治理要求。

无论您是独立的数据科学家还是大型 ML 平台团队的一员，模型注册表都是扩展可靠且可维护的机器学习系统的基础组件。

概念​

模型注册表引入了一些概念，用于描述和促进 MLflow 模型的完整生命周期。

概念描述
模型

MLflow 模型是使用模型风味的 mlflow..log_model() 方法或从 MLflow 3 开始的 mlflow.create_external_model() API 创建的。一旦记录，此模型就可以注册到模型注册表。

注册模型

MLflow 模型可以注册到模型注册表。注册模型具有唯一的名称，包含版本、别名、标签和其他元数据。

模型版本

每个注册的模型可以有一个或多个版本。当新模型添加到模型注册表时，它将作为版本 1 添加。注册到相同模型名称的每个新模型都会递增版本号。模型版本具有标签，这对于跟踪模型版本的属性很有用（例如 pre_deploy_checks: "PASSED"）

模型 URI

您可以使用以下格式的 URI 引用注册模型：models://，例如，如果您有一个名为“MyModel”且版本为 1 的注册模型，则引用该模型的 URI 是：models:/MyModel/1”。

模型别名

模型别名允许您为注册模型的特定版本分配一个可变的命名引用。通过将别名分配给特定的模型版本，您可以通过模型 URI 或模型注册表 API 使用该别名引用该模型版本。例如，您可以创建一个名为 champion 的别名，指向名为 MyModel 的模型的版本 1。然后，您可以使用 URI models:/MyModel@champion 引用 MyModel 的版本 1。

别名对于部署模型特别有用。例如，您可以将 champion 别名分配给用于生产流量的模型版本，并在生产工作负载中以此别名为目标。然后，您可以通过将 champion 别名重新分配给不同的模型版本来更新服务于生产流量的模型。

标签

标签是与注册模型和模型版本关联的键值对，允许您按功能或状态对其进行标记和分类。例如，您可以将一个键为 "task" 值为 "question-answering"（在 UI 中显示为 task:question-answering）的标签应用于用于问答任务的注册模型。在模型版本级别，您可以将正在进行部署前验证的版本标记为 validation_status:pending，将已清除可部署的版本标记为 validation_status:approved。

注释和描述

您可以使用 Markdown 单独注释顶级模型和每个版本，包括描述和任何对团队有用的相关信息，例如算法描述、使用的数据集或给定版本建模方法所涉及的整体方法。

模型注册表实践​

MLflow 模型注册表在开源 (OSS) MLflow 和 Databricks 等托管平台中均可用。根据环境的不同，注册表提供不同级别的集成、治理和协作功能。

OSS MLflow 中的模型注册表​

在 MLflow 的开源版本中，模型注册表提供 UI 和 API 来管理机器学习模型的生命周期。您可以注册模型、跟踪版本、添加标签和描述，以及在预生产和生产等阶段之间转换模型。

在 MLflow 中注册模型

使用 MLflow Python API 注册模型​

MLflow 提供了几种注册模型版本的方法

# Option 1: specify `registered_model_name` parameter when logging a model
mlflow.<flavor>.log_model(..., registered_model_name="<YOUR_MODEL_NAME>")

# Option 2: register a logged model
mlflow.register_model(model_uri="<YOUR_MODEL_URI>", name="<YOUR_MODEL_NAME>")

注册模型后，您可以使用模型名称和版本将其加载回来

mlflow.<flavor>.load_model("models:/<YOUR_MODEL_NAME>/<YOUR_MODEL_VERSION>")

要了解有关 OSS 模型注册表的更多信息，请参阅模型注册表教程。

Databricks 中的模型注册表​

Databricks 通过将模型注册表与 Unity Catalog 集成来扩展 MLflow 的功能，从而实现集中治理、细粒度访问控制和跨工作区协作。

Unity Catalog 集成的关键优势包括：

  • 🛡️ 增强治理：对模型资产应用访问策略和权限控制。
  • 🌐 跨工作区访问：注册一次模型即可在多个 Databricks 工作区中访问。
  • 🔗 模型血缘：跟踪用于创建每个模型的笔记本、数据集和实验。
  • 🔍 发现和重用：从共享目录中浏览和重用生产级模型。
在 Databricks UC 中注册模型

使用 MLflow Python API 将模型注册到 Databricks UC​

先决条件：将跟踪 URI 设置为 Databricks

import mlflow

mlflow.set_registry_uri("databricks-uc")

使用 MLflow API 注册模型

# Option 1: specify `registered_model_name` parameter when logging a model
mlflow.<flavor>.log_model(..., registered_model_name="<YOUR_MODEL_NAME>")

# Option 2: register a logged model
mlflow.register_model(model_uri="<YOUR_MODEL_URI>", name="<YOUR_MODEL_NAME>")
警告

UC 中的 ML 模型版本必须具有模型签名。如果您想在已记录或保存的模型上设置签名，可以使用mlflow.models.set_signature() API。

注册模型后，您可以使用模型名称和版本将其加载回来

mlflow.<flavor>.load_model("models:/<YOUR_MODEL_NAME>/<YOUR_MODEL_VERSION>")

有关更多信息，请参阅有关管理模型生命周期的 Databricks 文档。