跳到主要内容

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

·阅读时长12分钟
Awadelrahman M. A. Ahmed
Awadelrahman M. A. Ahmed
MLflow 大使 | REMA 1000 云数据与分析架构师

我们都（好吧，我们中的大多数人）还记得2022年11月，OpenAI 公开推出 ChatGPT，这标志着人工智能世界的一个重要转折点。虽然生成式人工智能 (GenAI) 已经发展了一段时间，但基于 OpenAI 的 GPT-3.5 架构构建的 ChatGPT 迅速抓住了公众的想象力。这导致 GenAI 在科技行业和公众中引起了爆炸性的兴趣。

在工具方面，MLflow 继续巩固其作为 ML 社区中（机器学习操作）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 这样的框架允许开发者构建将各种 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。如果您不熟悉它，可以将其视为 MLflow 中一个通用接口，通过定义一个自定义 Python 函数，将任何框架中的任何模型转换为 MLflow 模型。如果您希望更深入地了解，也可以参考 这篇早期文章。

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

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

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

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

    📌 但是等等！重要提示：

    • 我们的 model_code.py 脚本需要调用（即包含）mlflow.models.set_model() 来设置模型，这对于使用 load_model() 进行推理时重新加载模型至关重要。您会在示例中注意到这一点。

  2. 第二个文件记录了我们（在 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 文件：​

这也可以在 Jupyter Notebook 中定义。以下是其主要内容：

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 示例仅仅是一个开始——还有更多值得探索。下面列出了一些您可能会觉得有用的代码示例：

  1. 使用 Pyfunc log model API 记录代码模型 ( 模型代码 | 驱动代码 )
  2. 使用 Langchain log model API 记录代码模型 ( 模型代码 | 驱动代码 )