训练机器学习模型通常一开始是有条不紊的,但最终往往会陷入一片混乱之中。

我们都经历过这种情况:数十个实验分散在各种随机文件夹中,模型文件也被保存成像model_v2_final_FINAL.pkl这样的名字,因为没人能确定到底哪个版本是有效的。

一旦你从个人项目转向团队合作,或者试图将某些成果投入实际生产环境,这种“有组织的混乱”很快就会成为严重的瓶颈。

要解决这个问题,仅仅依靠更好的命名规范是不够的;我们还需要一种标准化的方法来记录和交接工作。而MLflow正是为填补这一空白而开发的。

MLflow最初由Databricks团队在2018年推出,如今它已经成为了管理整个机器学习生命周期的标准开源平台。它充当了一个中央枢纽,让你的实验、代码和模型都集中在这里,而不会被遗忘在各个文件夹里。

在本教程中,我们将介绍MLflow的核心理念,以及它的模块化架构是如何解决机器学习领域“依赖关系混乱”这一问题的。我们会详细讲解跟踪系统、项目管理、模型构建以及模型注册表这四个关键组成部分,并通过实际操作来帮助你将项目从本地开发环境顺利迁移到可投入生产的环境。

目录:

先决条件:

为了充分利用本教程的内容,你应当具备以下条件:

  • 基本的Python编程能力:熟悉上下文管理器(with语句)和装饰器的使用。

  • 机器学习基础知识:了解训练/测试数据的划分方式,以及模型评估指标(如准确率、损失值等)的含义。

  • 本地开发环境:已安装Python 3.8或更高版本。熟悉使用pipconda来安装软件包会很有帮助。

MLflow架构:整体概览

要理解为什么MLflow如此高效,就必须了解它的构成方式。MLflow并非一个庞大而僵化的工具,而是一个由四个相互独立的核心组件构成的模块化系统。

这一设计非常重要,因为这意味着你不必一次性使用整个生态系统。如果你只需要跟踪实验过程,而对其他功能并不感兴趣,那么只需使用相关部分即可,无需关注其余内容。

为了让大家更直观地理解这些组件的作用,下面将它们与你们可能已经熟悉的一些工具进行对比:

  • MLflow Tracking:用于记录实验过程、各项指标以及参数信息。(可以将其理解为“用于跟踪ML训练流程的Git提交记录”)

  • MLflow Projects:用于打包代码以确保其可重复执行性。(相当于“用于存储ML代码的Docker镜像”)

  • MLflow Models:为多种编程框架提供统一的模型存储格式。(可以理解为“适用于不同框架的通用模型适配器”)

  • Model Registry:负责管理模型的版本控制及相关信息。(相当于“用于模型管理的CI/CD流程”)

从架构角度来看,MLflow主要由客户端和服务器两部分组成。

你大部分时间都在使用客户端——无论是训练脚本还是Jupyter笔记本,都是通过这些工具来记录指标数据或注册模型的。

而服务器则充当后台处理系统,负责数据的存储工作。它主要包括跟踪服务器、后端存储系统(通常为PostgreSQL等数据库)以及模型文件存储服务(如S3或GCS)。这些大型存储资源用于保存模型权重等重要文件。

正是这种分离机制使得MLflow具备极高的灵活性。最初,你可以在笔记本电脑上仅使用文件系统即可开始使用MLflow;当团队规模扩大时,只需将本地环境替换为集中式服务器和云存储服务,而无需对代码进行任何修改。这样的设计让MLflow能够随着项目的发展而逐步扩展,而不会在项目进入关键阶段时迫使你重新开始。

现在,让我们来详细了解一下MLflow的这四个核心组件,以便大家更好地理解它们的工作原理。

了解MLflow Tracking机制

对大多数团队来说,跟踪功能是使用MLflow的第一步。它的作用非常简单:它就像一本数字实验记录本,会记录下训练过程中发生的所有信息。

你无需费力去回忆学习率的具体数值或准确率图保存的位置,MLflow会自动在后台完成这些记录工作。

在这里,最基本的单位是训练运行。可以将一次训练运行视为你的训练代码的一次执行过程。在每次运行过程中,系统会捕获四种类型的特定信息:

  • 参数:输入数据,例如批量大小或模型中的树节点数量。

  • 指标:输出结果,比如准确率或损失值,这些数据可以随时间进行跟踪分析。

  • 模型文件:包括模型权重、混淆矩阵或图像等重要数据文件。

  • 标签与元数据:用于记录哪些开发人员执行了代码,以及使用了哪个Git提交版本等信息。

一个跟踪示例

在实际操作中看到这些流程,才是理解该架构真正工作原理的最佳方式。你无需重新构建整个开发流程——只需将训练逻辑封装在上下文管理器中即可。

以下是Python中基本集成方式的示例:

import mlflow 
import mlflow.sklearn 
from sklearn.ensemble import RandomForestClassifier 
from sklearn.metrics import accuracy_score 

# 这段代码用于启动运行并整理相关数据
with mlflow.start_run():    
    # 记录参数信息    
    mlflow.log_param("n_estimators", 100)    
    mlflow.log_param("max_depth", 5)    

    # 训练模型    
    model = RandomForestClassifier(n_estimators=100, max_depth=5)    
    model.fit(X_train, y_train)    

    # 记录评估指标    
    accuracy = accuracy_score(y_test, model.predict(X_test))    
    mlflow.log_metric("accuracy", accuracy)    

    # 将模型保存为资源文件    
    mlflow.sklearn.log_model(model, "random_forest_model")

MLflow用户界面中显示的对比表格,将三次训练过程并排展示,突出参数和指标的差异。

`mlflow.start_run()`这个上下文管理器会创建一个新的训练任务,并在该代码块执行完毕后自动关闭它。在该代码块内记录的所有数据都会与该训练任务关联起来,并存储在后端存储系统中。

数据实际上被保存在哪里?

当你在笔记本电脑上刚开始使用MLflow时,系统会自动创建一个本地的`./mlruns`目录来存储相关数据。但当你进入团队环境,并将所有人指向一个集中的跟踪服务器时,MLflow的强大功能才会真正显现出来。

系统会根据数据的规模来决定其存储方式:结构化数据(如参数和评估指标)体积较小,且需要便于查询,因此会被保存到PostgreSQL这样的SQL数据库中;而非结构化数据(比如模型文件或大型图表文件),由于体积较大,不适合存储在数据库中,因此会被保存到Amazon S3或Google Cloud Storage这样的资源存储系统中。

MLflow资源存储系统的界面截图,展示了已记录模型的目录结构,包括MLmodel元数据和.model.pkl文件。

为什么要费心设置这些流程呢?

如果依赖直觉或混乱的命名规则,那么当项目规模扩大后,肯定会遇到各种问题。这种做法可能在一两天内还能奏效,但一旦需要比较模型的二十个不同版本时,这种方式就会立刻暴露出其缺陷来。

通过将跟踪功能独立成一个独立的架构模块,MLflow能够为你提供可供查询的历史数据。你无需再翻阅那些旧笔记,只需进入用户界面进行筛选,就能快速找到最理想的结果,并清楚地了解是哪种配置方案使你取得了想要的效果。这样一来,数据科学中的“推理”环节就不再需要依靠猜测来完成啦。

这是一张MLflow并行坐标图,用于展示在多次运行过程中,估计器数量与模型精度之间的关系。
这是一张MLflow散点图,用于说明n_estimators参数与模型精度之间的正相关关系。

了解MLflow项目

即使你能训练出世界上最精确的模型,但如果你的同事在他们的机器上无法重现你的实验结果,那么这个模型也就没有什么实际价值了。

这时MLflow项目就派上了用场。它们通过提供一种标准化的方法来打包你的代码、依赖库以及执行入口点,从而解决了可重复性带来的问题。

可以把MLflow项目想象成一个目录(或者一个Git仓库),在其根目录下有一个名为MLproject的特殊“使用说明书”文件。这个文件会明确告诉任何人或任何服务器需要什么样的运行环境,以及如何启动代码的执行过程。

MLproject文件

你不需要向别人发送一份冗长的安装步骤说明文档,只需将MLproject文件交给他们即可。以下是一个典型的训练流程配置示例:

name: my_ml_project
conda_env: conda.yaml

entry_points:
  train:
    parameters:
      learning_rate: {type: float, default: 0.01}
      epochs: {type: int, default: 50}
      data_path: {type: str}
    command: "python train.py --lr {learning_rate} --epochs {epochs} --data {data_path}"
  
  evaluate:
    parameters:
      model_path: {type: str}
    command: "python evaluate.py --model {model_path}"

conda_env这一行指定了一个conda.yaml文件,其中列出了你的代码所需的Python包及其版本。如果你需要更高的隔离性,MLflow也支持Docker环境。

这种配置方式的优点在于它的简洁性。任何安装了MLflow的人都可以通过一条命令来运行你的整个项目:

mlflow run . -P learning_rate=0.001 -P epochs=100 -P data_path=./data/train.csv

为什么这很重要

在两种特定场景下,MLflow项目的作用尤为显著。首先是新成员的快速上手过程:新加入团队的成员可以克隆你的仓库,并在几分钟内就开始运行代码,而无需花费一整天的时间去解决库版本冲突的问题。

其次是持续集成与持续部署流程。由于这些项目是可以被程序化地触发的,因此它们非常适合融入自动化训练流程中。当可重复性成为至关重要的因素时,拥有一个统一的、标准化的代码执行方式,会让所有相关人员的工作都变得更加顺利。

了解MLflow模型注册系统

通过跟踪实验可以确定哪个模型是“获胜者”,但真正管理这个模型从开发环境到生产环境的整个流程的地方,就是模型注册系统。

可以把模型注册系统看作是一种治理机制。它负责处理模型的版本控制、阶段管理,并生成清晰的审计记录,这样你就永远不会不知道当前在 production 环境中运行的是哪个模型。

模型注册系统使用一些简单的概念来保持各项信息的条理清晰:

  • 已注册的模型:这就是你项目的整体名称,比如 “CustomerChurnPredictor”。

  • 模型版本:每次你推送一个新的模型版本时,MLflow 会自动为该版本分配一个编号(如 v1、v2 等)。

  • 阶段:这些标签用于表示模型的当前所处的生命周期阶段,比如 “测试阶段”、“生产环境” 或 “已归档”。

  • 注释:这些只是用来记录一些信息的备注或标签。它们有助于说明为什么某个特定版本会被启用,或者该版本存在哪些特殊问题。

MLflow模型注册系统的界面显示:IrisClassifier模型的第1个版本已正式转入生产环境。

将模型推送到相应的阶段

在现实的工作流程中,你并不会简单地“部署”一个文件,而是会将其逐步推进到不同的开发或测试阶段。使用 MLflow Client 来操作的话,具体步骤如下:

Python
import mlflow
from mlflow.tracking import MlflowClient

client = MlflowClient()

# 首先,我们将一次成功的实验中生成的模型注册到系统中
result = mlflow.register_model(
    model_uri=f"runs:/{run_id}/random_forest_model",
    name="CustomerChurnPredictor"
)

# 然后,我们将第1个版本模型推进到测试阶段,以便质量保障团队进行审查
client.transition_model_version_stage(
    name="CustomerChurnPredictor",
    version=1,
    stage="Staging"
)

# 当一切检查都通过后,我们将其推送到生产环境
clienttransition_model_version_stage(
    name="CustomerChurnPredictor",
    version=1,
    stage="Production"
)

为什么这很重要?

模型注册系统解决了一个问题:当团队规模扩大时,往往很难准确掌握哪些版本已经在生产环境中运行、是谁批准了这些版本的部署,以及这些版本是相对于什么标准进行评估的。如果没有这样的系统,这些信息很可能会分散在各种聊天记录或过时的电子表格中,导致信息混乱。

此外,模型注册系统还使得回滚操作变得非常简单。如果生产环境中的第3个版本出现了问题,你不需要重新部署整个系统,只需将第2个版本重新推回到注册系统中的“生产阶段”即可。由于你的服务基础设施会自动选择使用“生产阶段”的模型版本,因此系统会自动切换回稳定的版本。

各组件是如何协同工作的

要想了解这些机制在现实世界中究竟如何运作,不妨从头到尾仔细梳理一下典型的工作流程。其实,这个过程就像一场接力赛,每个组件都会将“接力棒”传递给下一个组件。

整个流程始于数据科学家进行一系列实验。每当他们启动实验时,MLflow Tracking系统就会在后台默默记录各项数据,自动统计相关指标,并将模型生成的各类文件保存到后端存储系统中。在这个阶段,所有的工作都围绕着探索和寻找最佳方案来进行。

一旦确定了最优的实验结果,该模型就会被正式注册到模型注册库中。这时团队就可以开始介入了:他们可以通过用户界面查看实验注释、审查评估结果,然后将模型移入测试阶段。在通过更多验证测试后,该模型就能被正式投入生产环境。

当真正需要部署模型时,部署系统会从注册库中获取当前的生产版本。无论你是使用Kubernetes、云端服务,还是MLflow内置的服务器,这一过程都是相同的。

由于MLproject文件负责处理各种依赖关系,而MLflow Models格式则负责管理框架相关细节,因此部署基础设施根本不需要关心模型是用Scikit-learn还是PyTorch构建的。所有的必要信息都已经准备齐全,所以交接过程非常顺利。

正是这样的工作流程,使得MLflow从一个由众多实用工具组成的集合,发展成了一个功能完备的机器学习运维平台。它将数据科学中那些繁琐的实验阶段,与生产环境中的严格规范紧密地联系在了一起。

总结

归根结底,MLflow的设计初衷就是尽量避免给用户带来额外的麻烦。它不会强迫你改变编写代码的方式或使用的库,而是提供必要的结构,帮助你的机器学习项目具备可复现性,并便于团队进行协作管理。

无论你是只是想避免给文件起名为“model_final_v2.pkl”,还是想要为模型构建复杂的持续集成/持续部署流程,了解这四个核心要素都是最好的起点。最有效的学习方法就是立即启动一个本地的跟踪服务器开始记录数据。你会发现,一旦拥有了这些实验数据的“真实记录”,你就再也不会想回到以前的工作方式了。

Comments are closed.