MLflow 中的模型代码日志记录 - 是什么、为什么以及如何

2024年9月13日 ·13 分钟阅读

MLflow 大使 | REMA 1000 的云数据与分析架构师

我们所有人都（好吧，大多数人）还记得 2022 年 11 月 OpenAI 公开发布 ChatGPT，这标志着人工智能世界的一个重要转折点。尽管生成式人工智能（GenAI）已经发展了一段时间，但基于 OpenAI GPT-3.5 架构的 ChatGPT 迅速抓住了公众的想象力。这引发了技术行业和广大公众对 GenAI 爆炸式的兴趣。

在工具方面，MLflow 继续巩固其作为 ML 社区中最受欢迎的（机器学习运维）MLOps 工具的地位。然而，GenAI 的兴起带来了我们在使用 MLflow 方面的新需求。这些新挑战之一就是如何在 MLflow 中记录模型。如果你以前使用过 MLflow（我敢打赌你用过），你可能很熟悉 mlflow.log_model() 函数及其如何高效地 pickle 模型工件。

特别是对于 GenAI，出现了一个新要求：将模型“从代码中”记录下来，而不是将其序列化为 pickle 文件！你猜怎么着？这个需求不仅限于 GenAI 模型！因此，在本文中，我将探讨这个概念以及 MLflow 如何适应以满足这一新要求。

你会注意到这项功能是在一个非常抽象的层面上实现的，允许你“将任何模型记录为代码”，无论是 GenAI 还是其他模型！我喜欢将其视为一种通用方法，其中 GenAI 模型只是其用例之一。因此，在本文中，我将探讨这项新功能——“从代码记录模型”。

阅读完本文后，你应该能够回答三个主要问题：“是什么”、“为什么”以及“如何”使用从代码记录模型。

什么是从代码记录模型？

事实上，当 MLflow 宣布这项功能时，它让我对“模型”这个概念进行了更抽象的思考！如果你退后一步，将模型视为描述输入和输出变量之间关系的数学表示或函数，你可能会觉得这很有趣。在这个抽象层面上，一个模型可以是很多东西！

人们甚至可能意识到，一个模型作为一个对象或工件，仅仅代表了模型可以成为的东西的一种形式，即使它是 ML 社区中最流行的一种。如果你仔细想想，一个模型也可以很简单，比如一段用于映射函数的代码，或者一段向外部服务（如 OpenAI 的 API）发送 API 请求的代码。

我稍后会在文章中详细解释如何从代码记录模型的具体工作流程，但就目前而言，让我们从高层次来看，它包含两个主要步骤：首先，编写模型代码；其次，从代码记录模型。这如下面的图所示：

从代码记录模型的高层工作流程：

High Level Models-from-Code Logging Workflow

🔴 需要注意的是，当我们提到“模型代码”时，我们指的是可以被视为模型本身的**代码**。这意味着它**不是**生成已训练模型对象的训练代码，而是作为模型本身执行的逐步代码。

从代码记录模型与基于对象记录有何不同？

在前一节中，我们讨论了从代码记录模型的概念。然而，概念在与替代方案进行对比时通常会更清晰；这是一种称为对比学习的技术。在我们的案例中，替代方案是基于对象记录，这是 MLflow 中通常用于记录模型的常用方法。

基于对象记录将训练好的模型视为一个可以存储和重用的对象。训练后，模型被保存为一个对象，可以轻松加载以进行部署。例如，这个过程可以通过调用 mlflow.log_model() 来启动，MLflow 负责序列化，通常使用 Pickle 或类似方法。

基于对象记录可以分解为三个高层步骤，如下面的图所示：首先，创建模型对象（通过训练或获取它）；其次，对其进行序列化（通常使用 Pickle 或类似工具）；最后，将其记录为对象。

基于对象记录的高层工作流程：

High Level Object-Based Logging Workflow

💡 最流行的基于对象记录与从代码记录模型的主要区别在于，在前者中，我们记录的是模型对象本身，无论是你训练的模型还是你获取的预训练模型。而在后者中，我们记录的是代表你模型的代码。

何时需要从代码记录模型？

到目前为止，我希望你对什么是“从代码记录模型”有了清晰的认识！你可能仍然想知道该功能可以应用的具体用例。本节将确切地介绍这一点——“为什么”！

虽然我们在引言中提到了 GenAI 作为一个激励用例，但我们也强调了 MLflow 对从代码记录模型采取了更通用、更广泛的方法，我们将在下一节中看到这一点。这意味着你可以利用从代码记录模型的通用性来应对各种场景。我确定了三种我认为特别相关的关键使用模式：

1️⃣ 当你的模型依赖于外部服务时：

这是最明显和常见的用例之一，尤其是在现代人工智能应用兴起之际。我们越来越清楚地认识到，我们正在从“模型”粒度构建 AI 转向“系统”粒度构建 AI。

换句话说，人工智能不再仅仅关乎单个模型；它关乎这些模型如何在更广泛的生态系统中相互作用。随着我们越来越依赖外部 AI 服务和 API，从代码记录模型的需求变得更加突出。

例如，像 LangChain 这样的框架允许开发人员构建应用程序，将各种 AI 模型和服务链接起来以执行复杂任务，例如语言理解和信息检索。在这种情况下，“模型”不仅仅是一组可以被pickle 的训练参数，而是一个相互连接的“系统”，通常由发出对外部平台 API 调用的代码来编排。

在这种情况下，从代码记录模型可确保保留整个工作流程，包括逻辑和依赖项。它提供了通过捕获代码来保持相同模型体验的能力，从而能够忠实地重建模型的行为，即使实际的计算工作是在你的领域之外执行的。

2️⃣ 当你组合多个模型来计算复杂指标时：

除了 GenAI 之外，你仍然可以在各种其他领域受益于从代码记录模型功能。在许多情况下，多个专业模型被组合以产生全面的输出。请注意，我们不仅仅是指传统的集成建模（预测相同的变量）；通常，你需要组合多个模型来预测复杂推断任务的不同组成部分。

一个具体的例子可能是客户分析中的客户终身价值 (CLV)。在 CLV 的背景下，你可能有单独的模型用于

客户留存：预测客户将继续参与业务多长时间。
购买频率：预测客户购买的频率。
平均订单价值：估计每次交易的典型价值。

所有这些模型可能已经使用 MLflow 正确记录和跟踪。现在，你需要将这些模型“组合”成一个计算 CLV 的单一“系统”。我们称之为“系统”，因为它包含多个组件。

MLflow 从代码记录模型的优点在于，它允许你将这个“CLV 系统”视为一个“CLV 模型”。它使你能够利用 MLflow 的功能，维护类似 MLflow 的模型结构，并享受将 CLV 模型作为一个内聚单元进行跟踪、版本控制和部署的所有优势，即使它是建立在其他模型之上的。虽然可以使用自定义 MLflow PythonModel 构建这样一个复杂的模型系统，但利用从代码记录模型功能可以大大简化序列化过程，从而减少构建解决方案的阻力。

3️⃣ 当你根本没有序列化时：

尽管深度学习兴起，但行业仍然依赖于不产生序列化模型的基于规则的算法。在这些情况下，从代码记录模型对于将这些流程集成到 MLflow 生态系统中可能是有益的。

一个例子是在工业质量控制中，其中 Canny 边缘检测算法通常用于识别缺陷。这种基于规则的算法不涉及序列化，而是由特定步骤定义的。

另一个如今越来越受关注的例子是因果 AI。像 PC（Peter-Clark）算法这样的约束驱动因果发现算法可以发现数据中的因果关系，但它们是作为代码而不是模型对象实现的。

在任何一种情况下，通过从代码记录模型功能，你可以将整个过程作为一个“模型”记录在 MLflow 中，保留逻辑和参数，同时受益于 MLflow 的跟踪和版本控制功能。

如何实现从代码记录模型？

我希望到目前为止，你已经对从代码记录模型的“是什么”和“为什么”有了清晰的了解，现在你可能渴望亲自动手，专注于“如何做”！

在本节中，我将提供一个实现 MLflow 从代码记录模型的通用工作流程，然后是一个基本但广泛适用的示例。我希望这个工作流程能提供一个广泛的理解，使你能够应对各种场景。我也会在最后提供链接，指向涵盖更多特定用例（例如 AI 模型）的资源。

从代码记录模型工作流程：

实现的一个关键“配料”是 MLflow 的 pyfunc 组件。如果你不熟悉它，可以把 pyfunc 看作是 MLflow 中的一个通用接口，它允许你通过定义一个自定义 Python 函数，将任何框架的模型转换为 MLflow 模型。如果你想深入了解，也可以参考这篇早期的文章。

对于我们的从代码记录模型，我们将特别使用 pyfunc 中的 PythonModel 类。MLflow Python 客户端库中的此类允许我们创建和管理 Python 函数作为 MLflow 模型。它使我们能够定义一个自定义函数来处理输入数据并返回预测或结果。然后可以使用 MLflow 的功能来部署、跟踪和共享此模型。

这似乎正是我们所需要的——我们有一些代码充当我们的模型，并且我们想要记录它！这就是为什么你很快就会在我们的代码示例中看到 mlflow.pyfunc.PythonModel！

现在，每次我们需要实现从代码记录模型时，我们都会创建两个独立的 Python 文件：

第一个包含我们的模型代码（我们称之为 model_code.py）。该文件包含一个继承自 mlflow.pyfunc.PythonModel 类的类。我们定义的此类包含我们的模型逻辑。它可以是我们对 OpenAI API 的调用、CLV（客户终身价值）模型，或者我们的因果发现代码。我们很快就会看到一个非常简单的 101 示例。

📌 但等等！重要提示
- 我们的 model_code.py 脚本需要调用（即包含）mlflow.models.set_model() 来设置模型，这对于使用 load_model() 重新加载模型进行推理至关重要。你会在示例中注意到这一点。
第二个文件记录我们定义的类（在 model_code.py 中）。可以将其视为驱动代码；它可以是笔记本或 Python 脚本（我们称之为 driver.py）。在此文件中，我们将包含负责记录我们的模型代码（本质上是提供 model_code.py 的路径）的代码。

然后我们可以部署我们的模型。稍后，当加载服务环境时，将执行 model_code.py，当有服务请求进来时，将调用 PyFuncClass.predict()。

下图提供了这两个文件的通用模板。

Models from Code files

从代码记录模型的 101 示例：

让我们考虑一个简单的例子：一个基于直径计算圆面积的简单函数。通过从代码记录模型，我们可以将此计算记录为一个模型！我喜欢将其视为将计算框架化为一个预测问题，允许我们使用 predict 方法编写模型代码。

1. 我们的 `model_code.py` 文件：

import mlflow
import math

class CircleAreaModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        return [math.pi * (r ** 2) for r in model_input]

# It's important to call set_model() so it can be loaded for inference
# Also, note that it is set to an instance of the class, not the class itself.
mlflow.models.set_model(model=CircleAreaModel())

2. 我们的 `driver.py` 文件：

这也可以在笔记本中定义。以下是其基本内容

import mlflow

code_path = "model_code.py" # make sure that you put the correct path

with mlflow.start_run():
  logged_model_info = mlflow.pyfunc.log_model(
                                            python_model=code_path,
                                            artifact_path="test_code_logging"
                                            )

#We can proint some info about the logged model
print(f"MLflow Run: {logged_model_info.run_id}")
print(f"Model URI: {logged_model_info.model_uri}")

在 MLflow 中看起来是这样的：

执行 driver.py 将启动一个 MLflow 运行，并将我们的模型记录为代码。文件可以如以下所示：

Models from Code files

结论和后续学习

我希望到此为止，我已经履行了之前的承诺！你应该对什么是“从代码记录模型”及其与流行的将模型记录为序列化对象的基于对象记录方法的不同有了更清晰的认识。你也应该对为什么以及何时使用它有了坚实的基础，以及对如何通过我们的通用示例来实现它有了了解。

正如我们在引言和整篇文章中提到的，在许多情况下，从代码记录模型可能是有益的。我们的 101 示例只是一个开始——还有更多内容有待探索。下面是一些你可能会觉得有用的代码示例列表：

使用 **Pyfunc** log model API 从代码记录模型（模型代码 | 驱动代码）
使用 **Langchain** log model API 从代码记录模型（模型代码 | 驱动代码）

GenAI 应用和代理

模型训练

GenAI 应用和代理

模型训练

MLflow 中的模型代码日志记录 - 是什么、为什么以及如何