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

September 13, 2024 · 12 min read

MLflow Ambassador | Cloud Data & Analytics Architect at REMA 1000

我们都（嗯，大多数人）记得2022年11月，OpenAI公开发布ChatGPT，这标志着人工智能领域的一个重要转折点。尽管生成式人工智能（GenAI）已经发展了一段时间，基于OpenAI的GPT-3.5架构的ChatGPT迅速激发了公众的想象力。这在科技行业和普通大众中都引发了对GenAI的极大兴趣。

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

你会注意到，这个功能是在一个非常抽象的层面上实现的，允许你将任何模型“以代码的形式”记录下来，无论它是否是 GenAI！我倾向于将其视为一种通用方法，GenAI 模型只是其用例之一。因此，在这篇文章中，我将探讨这个新功能， "Models from Code logging"。

在本文结束时，你应该能够回答关于使用 Models from Code logging 记录模型的三个主要问题：'什么'、'为什么' 和 '如何'。

事实上，当 MLflow 宣布这个功能时，它让我以更抽象的方式思考“model”这一概念！如果你把视野放宽，把模型看作描述输入与输出变量关系的数学表示或函数，你也可能会觉得这很有趣。在这个抽象层面上，模型可以是多种形式！

人们甚至可能会意识到，模型作为一个对象或工件，仅代表模型可能呈现的某一种形式，即便它在机器学习社区中是最流行的形式。如果你仔细想想，模型也可以简单到像一段用于映射函数的代码，或者像一段向外部服务（例如 OpenAI's APIs）发送 API 请求的代码。

我将在文章后面详细解释如何从代码中记录模型的工作流程，但现在，我们先从高层次考虑两个主要步骤：首先，编写模型代码；其次，从代码中记录模型。如下图所示：

来自代码日志记录工作流的高级模型:

High Level Models-from-Code Logging Workflow

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

在上一节中，我们讨论了 Models from Code logging 的概念。然而，当将概念与其替代方案进行对比时，概念通常会更清晰；这是一种称为 对比学习 的技术。在我们的情况下，替代方案是 Object-Based logging，它是 MLflow 中用于记录模型的常用方法。

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

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

高级基于对象的日志记录工作流程:

High Level Object-Based Logging Workflow

💡流行的 Object-Based logging 与 Models from Code logging 之间的主要区别在于，前者记录的是模型对象本身，无论是你训练的模型还是你获取的预训练模型；而后者记录的则是表示你模型的代码。

何时需要来自代码日志记录的模型？

到目前为止，我希望你已经清楚地理解了什么是 Models from Code logging！不过，你可能仍然想知道这个功能可以应用于哪些具体用例。本节将正好涵盖这一点——也就是为什么！

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

这是一个明显且常见的用例之一，尤其是随着现代AI应用的兴起。越来越清楚的是，我们正在从以"模型"为粒度构建AI转向以"系统"为粒度构建AI。

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

在这些情况下，从代码记录模型可以确保整个工作流（包括逻辑和依赖）得到保留。它通过捕获代码提供了保持相同模型体验的能力，使得即使实际计算工作在您的领域之外进行，也能忠实地重现模型的行为。

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

除了 GenAI，你仍然可以在其他各种领域中从 Models from Code 功能中受益。许多情况下需要将多个专业化模型组合起来以产生全面的输出。请注意，我们这里并不是只指传统的集成建模（预测相同的变量）；通常，你需要结合多个模型来预测复杂推断任务的不同组成部分。

客户留存：预测客户会继续与业务互动多长时间。
购买频率：预测客户多久会进行一次购买。
平均订单价值：估算每笔交易的典型价值。

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

MLflow 的 Models from Code 日志功能的妙处在于，它允许你将这个“CLV system”视为一个“CLV model”。它使你能够利用 MLflow 的能力，保持 MLflow 风格的模型结构，并拥有跟踪、版本控制和将你的 CLV 模型作为一个整体部署的所有优势，即使它是构建在其他模型之上。虽然这样一个复杂的模型系统可以使用自定义的 MLflow PythonModel 来构建，但使用 Models from Code 功能可以大大简化序列化过程，减少构建该解决方案的摩擦。

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

在这两种情况下，使用 Models from Code 功能，您可以将整个过程记录为 MLflow 中的“模型”，保留逻辑和参数，同时享受 MLflow 的跟踪和版本控制功能。

如何从代码日志实现模型？

我希望到目前为止，你已经清楚地理解了“源自代码的模型”的“是什么”和“为什么”，现在你可能渴望亲自动手并专注于如何！

来自代码工作流的模型：

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

这正是我们要找的——我们有一些代码作为我们的模型，并且想把它记录下来！这就是为什么你很快会在我们的代码示例中看到mlflow.pyfunc.PythonModel！

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

📌 但请注意！重要：
- 我们的 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

让我们考虑一个简单的示例：一个根据直径计算圆面积的简单函数。借助 Models from Code，我们可以将此计算记录为一个模型！我喜欢把它看作是将该计算构建为一个预测问题，从而允许我们使用带有 predict 方法的模型代码。

1. 我们的 `model_code.py` file :

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())

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

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}")

Models from Code files

结论和进一步学习

我希望到现在为止，我已经履行了之前的承诺！你现在应该更清楚地理解什么是 Models from Code，以及它与将模型作为序列化对象记录的流行基于对象的方法有何不同。你还应该对 为什么 以及何时使用它有坚实的基础，并且了解通过我们的通用示例如何实现它。

正如我们在介绍中以及整篇文章中提到的，许多用例中从代码生成的模型都很有用。我们的 101 示例只是开始——还有更多可供探索。下面列出了一些你可能会觉得有用的代码示例：

从代码中使用 Pyfunc log model API 记录模型 ( model code | driver code )
使用 Langchain log model API 从代码中记录模型 ( model code | driver code )

Gen AI

Model training