跳到主要内容

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

·阅读时间12分钟
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() 函数以及它如何高效地 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️⃣ 当您的模型依赖于外部服务时:

这是显而易见且常见的用例之一,尤其是在现代人工智能应用程序兴起的情况下。我们越来越清楚地看到,我们正在从“模型”粒度构建人工智能转向“系统”粒度。

换句话说,人工智能不再仅仅关乎个体模型;它关乎这些模型在更广泛的生态系统中如何互动。随着我们越来越依赖外部人工智能服务和 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 文件:

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

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 从代码记录模型(模型代码 | 驱动代码