跳到主要内容

MLflow 跟踪

MLflow Tracking 是一个 API 和 UI,用于在运行机器学习代码时记录参数、代码版本、指标和输出文件,并用于后续可视化结果。MLflow Tracking 提供 PythonRESTRJava API。

MLflow Tracking UI 的截图,显示了模型训练期间验证损失指标的图表。

快速入门

如果您以前没有使用过 MLflow Tracking,我们强烈建议您按照以下快速入门教程进行操作。

概念

运行

MLflow Tracking 是围绕“运行”(runs)的概念组织的,“运行”是数据科学代码的执行,例如,单个 python train.py 执行。每次运行都会记录元数据(有关您的运行的各种信息,如指标、参数、开始和结束时间)和构件(运行的输出文件,如模型权重、图像等)。

模型

模型代表在运行期间生成的训练好的机器学习构件。已记录的模型包含其自身的元数据和类似于运行的构件。

实验

实验将运行和模型分组到一个特定的任务中。您可以使用 CLI、API 或 UI 创建实验。MLflow API 和 UI 也允许您创建和搜索实验。有关如何将运行分组到实验中的更多详细信息,请参阅 将运行组织到实验中

跟踪运行

MLflow Tracking API 提供了一组用于跟踪运行的函数。例如,您可以调用 mlflow.start_run() 来启动一个新的运行,然后调用 日志函数,如 mlflow.log_param()mlflow.log_metric() 分别记录参数和指标。请访问 Tracking API 文档 以获取有关使用这些 API 的更多详细信息。

import mlflow

with mlflow.start_run():
mlflow.log_param("lr", 0.001)
# Your ml code
...
mlflow.log_metric("val_loss", val_loss)

或者,自动日志记录提供了一个超快速的设置来开始 MLflow 跟踪。此强大功能允许您在无需显式日志语句的情况下记录指标、参数和模型 - 您需要做的就是在训练代码之前调用 mlflow.autolog()。自动日志记录支持 Scikit-learn、XGBoost、PyTorch、Keras、Spark 等流行库。有关支持的库以及如何将自动日志记录 API 与每个库一起使用的详细信息,请参阅 自动日志记录文档

import mlflow

mlflow.autolog()

# Your training code...
注意

默认情况下,在没有任何特定服务器/数据库配置的情况下,MLflow Tracking 会将数据记录到本地 mlruns 目录。如果您想将运行记录到其他位置,例如远程数据库和云存储,以便与团队共享结果,请按照 设置 MLflow Tracking 环境 部分中的说明进行操作。

以编程方式搜索已记录的模型

MLflow 3 通过 mlflow.search_logged_models() 引入了强大的模型搜索功能。此 API 允许您使用类似 SQL 的语法,根据性能指标、参数和模型属性在不同实验中查找特定模型。

import mlflow

# Find high-performing models across experiments
top_models = mlflow.search_logged_models(
experiment_ids=["1", "2"],
filter_string="metrics.accuracy > 0.95 AND params.model_type = 'RandomForest'",
order_by=[{"field_name": "metrics.f1_score", "ascending": False}],
max_results=5,
)

# Get the best model for deployment
best_model = mlflow.search_logged_models(
experiment_ids=["1"],
filter_string="metrics.accuracy > 0.9",
max_results=1,
order_by=[{"field_name": "metrics.accuracy", "ascending": False}],
output_format="list",
)[0]

# Load the best model directly
loaded_model = mlflow.pyfunc.load_model(f"models:/{best_model.model_id}")

主要特性

  • 类似 SQL 的过滤:使用 metrics.params. 和属性前缀来构建复杂的查询
  • 数据集感知搜索:根据特定数据集过滤指标,以便进行公平的模型比较
  • 灵活排序:按多个条件排序以查找最佳模型
  • 直接模型加载:使用新的 models:/<model_id> URI 格式进行即时模型访问

有关全面的示例和高级搜索模式,请参阅 搜索已记录模型指南

以编程方式查询运行

您还可以使用 MlflowClient 以编程方式访问 Tracking UI 中的所有功能。

例如,以下代码段搜索在实验的所有运行中具有最佳验证损失的运行。

client = mlflow.tracking.MlflowClient()
experiment_id = "0"
best_run = client.search_runs(
experiment_id, order_by=["metrics.val_loss ASC"], max_results=1
)[0]
print(best_run.info)
# {'run_id': '...', 'metrics': {'val_loss': 0.123}, ...}

跟踪模型

MLflow 3 引入了增强的模型跟踪功能,允许您在单个运行中记录多个模型检查点,并跟踪它们在不同数据集上的性能。这对于深度学习工作流特别有用,您希望在不同的训练阶段保存和比较模型检查点。

记录模型检查点

您可以使用模型日志函数中的 step 参数在训练的不同阶段记录模型检查点。每个记录的模型都会获得一个唯一的模型 ID,您可以在以后引用它。

import mlflow
import mlflow.pytorch

with mlflow.start_run() as run:
for epoch in range(100):
# Train your model
train_model(model, epoch)

# Log model checkpoint every 10 epochs
if epoch % 10 == 0:
model_info = mlflow.pytorch.log_model(
pytorch_model=model,
name=f"checkpoint-epoch-{epoch}",
step=epoch,
input_example=sample_input,
)

# Log metrics linked to this specific model checkpoint
accuracy = evaluate_model(model, validation_data)
mlflow.log_metric(
key="accuracy",
value=accuracy,
step=epoch,
model_id=model_info.model_id, # Link metric to specific model
dataset=validation_dataset,
)

将指标链接到模型和数据集

MLflow 3 允许您将指标链接到特定的模型检查点和数据集,从而提高模型性能的可追溯性。

# Create a dataset reference
train_dataset = mlflow.data.from_pandas(train_df, name="training_data")

# Log metric with model and dataset links
mlflow.log_metric(
key="f1_score",
value=0.95,
step=epoch,
model_id=model_info.model_id, # Links to specific model checkpoint
dataset=train_dataset, # Links to specific dataset
)

搜索和排序模型检查点

使用 mlflow.search_logged_models() 根据其性能指标搜索和排序模型检查点。

# Search for all models in a run, ordered by accuracy
ranked_models = mlflow.search_logged_models(
filter_string=f"source_run_id='{run.info.run_id}'",
order_by=[{"field_name": "metrics.accuracy", "ascending": False}],
output_format="list",
)

# Get the best performing model
best_model = ranked_models[0]
print(f"Best model: {best_model.name}")
print(f"Accuracy: {best_model.metrics[0].value}")

# Load the best model for inference
loaded_model = mlflow.pyfunc.load_model(f"models:/{best_model.model_id}")

MLflow 3 中的模型 URI

MLflow 3 引入了一种新的模型 URI 格式,该格式使用模型 ID 而不是运行 ID,从而提供更直接的模型引用。

# New MLflow 3 model URI format
model_uri = f"models:/{model_info.model_id}"
loaded_model = mlflow.pyfunc.load_model(model_uri)

# This replaces the older run-based URI format:
# model_uri = f"runs:/{run_id}/model_path"

这种新方法提供了几个优点

  • 直接模型引用:无需了解运行 ID 和构件路径
  • 更好的模型生命周期管理:每个模型检查点都有自己的唯一标识符
  • 改进的模型比较:轻松比较同一运行中的不同检查点
  • 增强的可追溯性:模型、指标和数据集之间的清晰链接

跟踪数据集

MLflow 提供了跟踪与模型训练事件相关联的数据集的能力。数据集关联的元数据可以通过使用 mlflow.log_input() API 来存储。要了解更多信息,请访问 MLflow 数据文档,了解此 API 中提供的功能。

探索运行、模型和结果

Tracking UI

Tracking UI 可视化地展示您的实验、运行和模型,如本页顶部所示。

  • 基于实验的运行列表和比较(包括跨多个实验的运行比较)
  • 按参数或指标值搜索运行
  • 可视化运行指标
  • 下载运行结果(构件和元数据)

这些功能也适用于模型,如下所示。

MLflow UI Experiment view page models tab

MLflow Tracking UI 在模型选项卡上的截图,显示了实验下的模型列表。

如果您将运行记录到本地 mlruns 目录,请在它的上一级目录中运行以下命令,然后在浏览器中访问 http://127.0.0.1:5000

mlflow ui --port 5000

或者,MLflow Tracking Server 提供相同的 UI,并支持远程存储运行构件。在这种情况下,您可以从任何可以连接到您的跟踪服务器的机器上查看 UI,地址为 http://<您的 MLflow 跟踪服务器的 IP 地址>:5000

设置 MLflow Tracking 环境

注意

如果您只想将实验数据和模型记录到本地文件,可以跳过此部分。

MLflow Tracking 支持您开发工作流的多种不同场景。本节将指导您如何为您的特定用例设置 MLflow Tracking 环境。从宏观角度来看,MLflow Tracking 环境由以下组件组成。

组件

MLflow Tracking API

您可以在 ML 代码中调用 MLflow Tracking API 来记录运行并根据需要与 MLflow Tracking Server 进行通信。

后端存储

后端存储持久化每个运行的各种元数据,如运行 ID、开始和结束时间、参数、指标等。MLflow 支持两种类型的后端存储:**基于文件系统**(如本地文件)和**基于数据库**(如 PostgreSQL)。

此外,如果您正在与托管服务(如 Databricks 或 Azure Machine Learning)进行交互,您将与基于 REST 的后端存储进行交互,该存储由外部管理且无法直接访问。

构件存储

构件存储持久化(通常很大)的每次运行的构件,例如模型权重(如序列化的 scikit-learn 模型)、图像(如 PNG)以及模型和数据文件(如 Parquet 文件)。MLflow 默认将构件存储在本地文件 (mlruns) 中,但也支持不同的存储选项,如 Amazon S3 和 Azure Blob Storage。

对于作为 MLflow 构件记录的模型,您可以通过 `models:/` 格式的模型 URI 来引用该模型,其中“model_id”是分配给已记录模型的唯一标识符。这取代了旧的 `runs://` 格式,并提供了更直接的模型引用。

如果模型已在MLflow 模型注册表中注册,您还可以通过 `models://` 格式的模型 URI 来引用该模型,有关详细信息,请参阅MLflow 模型注册表

MLflow Tracking Server(可选)

MLflow Tracking Server 是一个独立的 HTTP 服务器,提供用于访问后端和/或构件存储的 REST API。Tracking Server 还提供了配置要提供的数据、管理访问控制、版本控制等的灵活性。请阅读 MLflow Tracking Server 文档 以获取更多详细信息。

常见设置

通过正确配置这些组件,您可以创建一个适合您团队开发工作流的 MLflow Tracking 环境。下图和表格展示了几种 MLflow Tracking 环境的常见设置。

1. 本地主机(默认)2. 本地数据库的本地跟踪3. 使用MLflow Tracking Server进行远程跟踪
场景独立开发独立开发团队开发
用例默认情况下,MLflow 将每次运行的元数据和构件记录到本地目录 `mlruns`。这是开始使用 MLflow Tracking 的最简单方法,无需设置任何外部服务器、数据库和存储。MLflow 客户端可以与 SQLAlchemy 兼容的数据库(例如 SQLite、PostgreSQL、MySQL)进行交互以实现后端。将元数据保存到数据库可以更清晰地管理您的实验数据,同时无需设置服务器。MLflow Tracking Server 可以配置为构件 HTTP 代理,通过跟踪服务器传递构件请求,以存储和检索构件,而无需与底层对象存储服务进行交互。这对于您希望在具有适当访问控制的共享位置存储构件和实验元数据的团队开发场景特别有用。
教程快速入门使用本地数据库跟踪实验使用 MLflow 追踪服务器进行远程实验追踪

使用MLflow Tracking Server进行其他配置

MLflow Tracking Server 为其他特殊用例提供了可定制性。请遵循 使用 MLflow Tracking Server 进行远程实验跟踪 来学习基本设置,并继续阅读以下材料以获取高级配置以满足您的需求。

在本地使用 MLflow Tracking Server

您当然可以在本地运行 MLflow Tracking Server。虽然这与直接使用本地文件或数据库相比没有太多额外的好处,但可能有助于在本地测试您的团队开发工作流或在容器环境中运行您的机器学习代码。

常见问题

我可以并行启动多个运行吗?

是的,MLflow 支持并行启动多个运行,例如多进程/多线程。有关更多详细信息,请参阅 在一个程序中启动多个运行

如何整齐地组织许多 MLflow 运行?

MLflow 提供了几种组织运行的方法

  • 将运行组织到实验中 - 实验是运行的逻辑容器。您可以使用 CLI、API 或 UI 创建实验。
  • 创建子运行 - 您可以在单个父运行下创建子运行以将它们分组。例如,您可以为交叉验证实验中的每个折叠创建一个子运行。
  • 为运行添加标签 - 您可以为每个运行关联任意标签,这允许您根据标签过滤和搜索运行。

我可以直接访问远程存储而不运行跟踪服务器吗?

是的,虽然建议将 MLflow Tracking Server 作为构件访问的代理以用于团队开发工作流,但如果您将其用于个人项目或测试,则可能不需要这样做。您可以通过遵循以下解决方法来实现此目的:

  1. 设置构件配置,例如凭据和端点,就像您为 MLflow Tracking Server 所做的那样。有关更多详细信息,请参阅 配置构件存储
  2. 创建一个具有显式构件位置的实验,
experiment_name = "your_experiment_name"
mlflow.create_experiment(experiment_name, artifact_location="s3://your-bucket")
mlflow.set_experiment(experiment_name)

此实验下的运行将直接将构件记录到远程存储。

如何将 MLflow Tracking 与模型注册表集成?

要将模型注册表功能与 MLflow 跟踪一起使用,您**必须使用基于数据库的存储**,例如 PostgresQL,并使用相应的模型类的 log_model 方法记录模型。一旦模型被记录,您就可以通过 UI 或 API 在模型注册表中添加、修改、更新或删除模型。有关如何为您的工作流正确配置后端存储,请参阅 后端存储常见设置

如何包含有关运行的附加描述性文本?

可以使用系统标签 mlflow.note.content 来添加此运行的描述性注释。虽然其他 系统标签 会自动设置,但此标签**默认不设置**,用户可以覆盖它以包含有关运行的附加信息。内容将在运行的页面上的“备注”部分显示。