跳到主要内容

MLflow 跟踪

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

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

快速入门

如果您以前没有使用过 MLflow Tracking,我们强烈建议您阅读以下快速入门教程。

概念

运行

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

模型

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

实验

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

跟踪运行

MLflow Tracking API 提供了一组函数来跟踪您的运行。例如,您可以调用 mlflow.start_run() 来开始一个新的运行,然后调用 日志函数,如 mlflow.log_param()mlflow.log_metric() 分别记录参数和指标。有关使用这些 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-learnXGBoostPyTorchKerasSpark 等。有关支持的库以及如何使用它们进行自动日志记录 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 中可用的功能。

探索运行、模型和结果

跟踪 UI

跟踪 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 并支持远程存储运行工件。在这种情况下,您可以从任何可以连接到您的跟踪服务器的机器访问 http://<您的 MLflow 跟踪服务器的 IP 地址>:5000 查看 UI。

设置 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 的后端存储进行接口,该存储无法直接访问。

工件存储

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

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

如果模型注册在 MLflow 模型注册表 中,您也可以通过 models:/<model-name>/<model-version> 格式的模型 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 创建实验。
  • 创建子运行 - 您可以在单个父运行下创建子运行以将其分组。例如,您可以为交叉验证实验中的每个折叠创建一个子运行。
  • 为运行添加标签 - 您可以为每个运行关联任意标签,这允许您根据标签筛选和搜索运行。

在不运行 Tracking Server 的情况下,我可以直接访问远程存储吗?

是的,虽然将 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 添加有关此运行的描述性说明。虽然其他系统标签是自动设置的,但此标签默认情况下未设置,用户可以覆盖它以包含有关运行的其他信息。内容将显示在运行页面的“备注”部分。