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

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

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

我们都（至少是大多数人）还记得 2022 年 11 月 OpenAI 公开推出 ChatGPT，这标志着人工智能世界的一个重要转折点。虽然生成式人工智能（GenAI）的发展已经有一段时间了，但 ChatGPT（基于 OpenAI 的 GPT-3.5 架构构建）迅速引起了公众的广泛关注。这导致了技术行业和公众对 GenAI 的兴趣爆炸式增长。

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

特别是对于 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️⃣ 当您的模型依赖于外部服务时：

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

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

例如，像 LangChain 这样的框架允许开发人员构建应用程序，将各种人工智能模型和服务链接起来执行复杂任务，例如语言理解和信息检索。在这种情况下，“模型”不仅仅是一组可以*序列化*的训练参数，而是一个由相互连接的服务组成的“系统”，通常由调用外部平台 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。如果您不熟悉它，可以将其视为 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 从代码记录模型（模型代码 | 驱动代码）

生成式 AI

模型训练

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