MLflow 模型注册表和 MinIO

Keith Pijanowski Keith Pijanowski 于 AI/ML 2023 年 8 月 11 日

介绍

MLflow 模型注册表允许您管理旨在用于生产环境的模型。这篇文章接续了我上一篇关于 MLflow 跟踪的文章。在我的跟踪文章中，我展示了如何记录参数、指标、工件和模型。如果您还没有阅读它，请在您有空的时候阅读一下。在这篇文章中，我将展示如何注册和版本化记录的模型。作为额外奖励，我还会展示如何将注册的模型从 MLflow 模型注册表加载到内存中。

MLflow 模型注册表是 MLflow 四个主要组件之一。下面简要描述了每个组件。

MLflow 跟踪 - 工程师将最常使用此功能。它允许记录和查询实验。它还跟踪每个实验的代码、数据、配置和结果。
MLflow 项目 - 允许通过将代码打包成平台无关的格式来重现实验。
MLflow 模型 - 将机器学习模型部署到可以提供服务的环境中。
MLflow 模型注册表 - 允许在中央存储库中存储、注释、发现和管理模型。

我假设您已经在您的开发机器上使用 PostgreSQL 和 MinIO 作为远程服务器安装了 MLflow。如果您还没有安装 MLflow，请查看我关于使用 MLflow 和 MinIO 设置开发机器的文章，我在文章中展示了如何使用 Docker Compose 运行下面显示的服务。

此部署中使用 MinIO 来存储模型和数据集。如果您要尝试模型并最终将它们移入生产环境，则拥有高质量的对象存储来存放这些工件至关重要。如果您正在训练大型语言模型，这一点尤其重要。

MLflow 和 MinIO 相辅相成。MinIO 是一个高性能对象存储（在 32 个 NVMe 节点上，GET 为 349 GB/s，PUT 为 177 GB/s），能够支持即使是最苛刻的数据湖、分析和 AI/ML 工作负载。数据以强一致性写入 MinIO，具有不可变对象。所有对象都通过内联擦除码、位腐烂散列和加密进行保护。

什么是注册模型？

注册模型是您生产环境的候选者。您应该在实验中注册性能最好的模型。注册模型后，您可以添加元数据，例如标签和描述。模型也可以在注册表中进行版本化，每个版本都可以有自己的描述和标签。此外，您可以指定模型的 `阶段` 来跟踪其状态。允许的阶段为 `None`、`Stage`、`Production` 和 `Archive`。我将在本文中演示所有这些功能。本文中显示的所有代码都可以在此处找到。

要注册模型，您必须之前使用我已在我的跟踪文章中演示的 `log_model()` 函数将其记录。回想一下，在这篇文章中，我使用了一个简单的 `log_model()` 函数变体将模型保存到 MLflow。下面显示了此行代码。

# 记录训练后的模型。
model_artifact_path = 'mnistmodel'
mlflow.pytorch.log_model(model, artifact_path=model_artifact_path)

为了利用模型注册表的全部功能，我们需要通过向其添加签名来升级此命令。签名基本上是模型的模式。它告诉 MLflow 允许用于推断的数据类型。您可以从下面修改的日志记录代码中看到，MLflow 提供了 `infer_signature()` 函数，顾名思义，该函数会根据您从训练集或测试集中发送的几个样本推断模型的签名。

# 记录训练后的模型。
images, labels = next(iter(test_loader))
img = images[0].view(1, 784)
label = labels[0]
print(img.size())
print(label.size())
with torch.no_grad():
logps = model(img)

model_signature = mlflow.models.infer_signature(img.numpy(), logps.numpy())
model_artifact_path = 'mnistmodel'
mlflow.pytorch.log_model(model, artifact_path=model_artifact_path,

signature=model_signature)

当我们有一个模型被记录并拥有一个模式后，我们就可以注册它。注册模型有两种方式。一种是使用 MLflow API。另一种是使用 MLflow UI 查找并注册先前记录的模型。我将展示这两种技术。让我们先从用于注册的 API 开始。

使用 API 注册模型

使用 API 注册先前记录的模型的最佳方法如下所示。还有其他两种技术，我将在稍后简要介绍，但以下技术是最好的，因为它允许你使用两个单独的 API 调用来创建顶级集合实体和模型的新版本。这使你有灵活性分别为它们提供不同的描述和不同的标签。可以将顶级集合实体视为一个对象，它维护一个包含相同名称的所有模型版本列表的对象。

model_name = 'mnist_nn_model'
client = MlflowClient()

# 如果顶级集合实体尚未注册，则注册它。
filter_string = f"name='{model_name}'"
results = client.search_registered_models(filter_string=filter_string)
if len(results) == 0:
model_tags = {'framework': 'Pytorch'}
model_description = '具有不同隐藏层的 MNIST 模型的各种版本。'
client.create_registered_model(model_name, model_tags, model_description)

# 注册新版本
run_id = active_run.info.run_id
run_uri = f'runs:/{run_id}/{model_artifact_path}'
model_source = RunsArtifactRepository.get_underlying_uri(run_uri)
version_tags = {'layers': len(params['hidden_sizes'])}
version_description = f'隐藏大小： {params["hidden_sizes"]}'
model_version = client.create_model_version(model_name, model_source, run_id,

tags=version_tags,

description=version_description)

`create_registered_model()` 函数创建集合实体。请注意，只有在集合不存在的情况下才能调用它，否则它会抛出错误。在上面的代码中，我使用了 `search_registered_models()` 函数来确保此函数不会对同一个模型调用两次。`create_model_version()` 函数创建模型的新版本。注意上面的代码中集合和版本具有不同的描述和标签。

注册模型的多个版本后，你将能够在 UI 中查看它们。让我们看一下 MLflow UI 中的模型注册表，并查看顶级条目和几个版本。通过导航到 http://localhost:5000/ 启动 UI。默认视图将显示你所有的实验。我们想要查看注册的模型，因此找到页面顶部的 `Models` 选项卡并单击它。你将看到一个已注册的所有模型的列表。下面显示了一个示例。

此视图显示所有注册的模型。对于每个模型，此列表显示注册的最新版本（版本 2）、已设置为暂存的最新版本（版本 1）以及已投入生产的最新版本（上面的屏幕截图中没有模型投入生产）。

单击模型名称将显示以下页面。

在这里，我们看到此模型的所有版本。此处显示的描述和标签适用于顶级条目。因此，它们应该表示所有版本的真实情况。

最后，MLflow 文档还描述了可用于注册模型的另外两组 API。针对每个框架专门开发的一些日志记录 API 具有注册功能。（这些将是在 mlflow.sklearn 中为 Sklearn 和 mlflow.pytorch 中为 Pytorch 找到的函数。）它们不提供与我在上面展示的相同级别的控制，并且对各个框架的支持也不同。还有一个 `mlflow.register_model()` 函数。此函数不提供对集合实体创建方式的控制，并且没有用于描述的参数。

使用 MLflow UI 注册模型

事实证明，你不需要 API 就可以注册模型。可以使用 MLflow UI 注册先前记录的模型。如果你希望人工确定应该将哪个记录的模型提升到注册表，这可能是注册模型的首选方法。让我们再次浏览 UI 以查看如何完成此操作。

首先，导航到 UI 的“实验”选项卡，并选择包含您要注册的模型的实验。

如果您将“模型”添加到每个运行显示的列列表中，那么您将能够看到哪些运行记录了模型，以及哪些记录的模型已注册。在上图中，已注册的模型具有名称和版本标识。选择一个包含您想要在模型注册表中的未注册模型的运行。这将带您到下面显示的运行详细信息页面。请注意，当您在此页面上的工件部分中选择模型时，将出现“注册模型”按钮。

单击“注册模型”按钮将显示下面的对话框。

在这里，如果我们需要新的集合实体，我们将创建顶级集合实体，或者我们可以将模型添加到已注册模型的现有集合中。完成此对话框后，您的模型将被注册。让我们导航到“模型”选项卡，并找到我们新注册的模型。

选择“mnist_lr_optimized”集合，您将看到此集合的版本列表。如下所示。另外请注意，集合级别上有一个描述和标签。

单击特定版本，我们可以查看该版本的详细信息。

该版本有自己的描述和标签。您可以使用此页面编辑这些值，也可以设置模型的阶段。还要查看此模型的架构。当我们加载模型并对其进行测试时，这将发挥作用，这将在下一节中介绍。

加载和测试已注册的模型

我们将做的最后一件事是加载以前注册的模型。毕竟，如果您无法检索模型并对其执行进一步的实验，那么记录和注册模型有什么意义？在本节中，我将展示的内容不应与称为“模型服务”的内容相混淆。当您服务模型时，您会采用以前训练和测试的模型，并将其与所有依赖项一起部署，使其能够作为独立服务执行推理。我们这里所做的是仅仅检索以前注册的模型。要使用此模型，您需要自己处理所有依赖项。如果模型仍在开发中，并且您想要查看同事刚刚训练的模型的新版本，这将非常有用。也许在训练模型时测试集不可用，并且您想看看新模型在从未见过的数据上的表现如何。

下面的代码将以前注册的模型加载到内存中。

model_name = 'mnist_nn_model'
model_version = 6

mlflow.set_tracking_uri('http://localhost:5000/')
model = mlflow.pyfunc.load_model(model_uri=f"models:/{model_name}/{model_version}")
type(model)

您只需要知道模型名称和版本号。这些值将用于创建模型 URI，该 URI 传递给 `mlflow.pyfunc.load_model()` 函数。上面的代码在笔记本单元格中运行时，将生成下面的输出。

mlflow.pyfunc.PyFuncModel

上面的代码中的类型检查告诉我们该模型是 `mlflow.pyfunc.PyFuncModel` 的实例。MLflow 使用 `PyFuncModel` 类包装所有模型。这似乎没有必要，但请记住，MLflow 是适用于所有框架的工具 - 因此，如果您在一个使用 Pytorch、TensorFlow 和 Sklearn 等多个框架进行实验的团队中，那么相同的下游测试或推理代码将适用于所有模型，因为它们都包装在一个通用类中。

我们包装的模型具有一个用于推理的 `predict()` 函数。让我们创建一个使用此模型的测试函数。

def test_model(model: mlflow.pyfunc.PyFuncModel, loader: DataLoader) -> Dict[str, Any]:
correct_count, total_count = 0, 0
for images, labels in loader:
for i in range(len(labels)):
img = images[i].view(1, 784)
# 关闭梯度以加快此部分的速度
with torch.no_grad():
logps = model.predict(img.numpy())

# 网络的输出是对数概率，

#需要对概率进行指数运算
ps = np.exp(logps)
probab = list(ps[0])
pred_label = probab.index(max(probab))
true_label = labels.numpy()[i]
if(true_label == pred_label):
correct_count += 1
total_count += 1

testing_metrics = {
'incorrect_count': total_count-correct_count,
'correct_count': correct_count,
'accuracy': (correct_count/total_count)
}
print("测试的图像数量 =", total_count)
print("\n模型准确率 =", (correct_count/total_count))
return testing_metrics

为了使上述函数能够工作，我们需要一个测试集。下面的代码将为我们提供MNIST数据集的测试集。

def load_test_images(batch_size: int) -> Tuple[Any]:
# 定义一个转换来规范化数据
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])

# 下载并加载测试数据
test_dataset = datasets.MNIST('./mnistdata', download=True, train=False,

transform=transform)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
return test_loader

这就是您需要从模型注册中心加载注册模型并对其进行进一步实验所需的一切。需要注意的是，如果包装模型的预测方法对您抛出错误，那么您的推断模式很可能出错。在这种情况下，请返回到记录模型的代码，并确保在 `infer_schema` 函数中使用的类型和形状与包装模型的预测方法接收的类型和形状相同。

摘要

将 MLflow 和 MinIO 相结合，可以产生强大的协同作用，提升机器学习开发流程。MLflow 简化的实验跟踪、模型版本控制和部署功能，确保数据科学家和工程师之间的有效协作和可重复性。通过集成 MinIO 高性能对象存储，该流程获得了可扩展性、成本效益和强大的数据管理功能，促进大规模数据集和模型的无缝存储和共享。这种统一的方法提高了团队效率，加速了创新，并为构建先进的机器学习应用程序奠定了坚实的基础。