MLflow 跟踪
MLflow Tracking 是一套 API 和 UI,用于在运行机器学习代码时记录参数、代码版本、指标和输出文件,并可在后续可视化这些结果。MLflow Tracking 提供 Python、REST、R 和 Java API。
快速入门
如果您之前从未使用过 MLflow Tracking,我们强烈建议您阅读以下快速入门教程。
概念
运行 (Runs)
MLflow Tracking 是围绕运行 (runs) 的概念组织的,运行即某一段数据科学代码的执行,例如单次 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)
或者,自动记录 (Auto-logging) 提供了一种启动 MLflow 追踪的超快速设置。这一强大的功能允许您无需显式编写记录语句即可记录指标、参数和模型——您所要做的就是在训练代码之前调用 mlflow.autolog()。自动记录支持 Scikit-learn、XGBoost、PyTorch、Keras、Spark 等主流库。请参阅自动记录文档以了解支持的库以及如何配合它们使用自动记录 API。
import mlflow
mlflow.autolog()
# Your training code...
默认情况下,如果没有进行特定的服务器/数据库配置,MLflow Tracking 会将数据记录到本地 mlruns 目录中。如果您希望将运行记录到其他位置(例如远程数据库或云存储),以便与团队共享结果,请按照设置 MLflow 追踪环境部分中的说明进行操作。
以编程方式搜索已记录的模型
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 以编程方式访问追踪 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)
# <RunInfo: run_id='...', experiment_id='0', status='FINISHED', start_time=...>
print(best_run.data.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 允许您直观地探索实验、运行和模型,如本页顶部所示。
- 基于实验的运行列表和比较(包括跨多个实验的运行比较)
- 按参数或指标值搜索运行
- 可视化运行指标
- 下载运行结果(制品和元数据)
这些功能也适用于模型,如下所示。
如果您将运行记录到本地 mlruns 目录,请在其父目录中运行以下命令,然后在浏览器中访问 http://127.0.0.1:5000。
mlflow server --port 5000
或者,MLflow 追踪服务器提供相同的 UI,并支持远程存储运行制品。在这种情况下,您可以从任何可以连接到追踪服务器的机器上通过 http://<MLflow 追踪服务器的 IP 地址>:5000 访问 UI。
设置 MLflow 追踪环境
如果您只想将实验数据和模型记录到本地文件,可以跳过此部分。
MLflow Tracking 支持您的开发工作流中的多种场景。本节将指导您如何为您的特定用例设置 MLflow 追踪环境。总体而言,MLflow 追踪环境由以下组件组成。
组件
MLflow Tracking API
您可以在 ML 代码中调用 MLflow Tracking API 来记录运行,并在必要时与 MLflow 追踪服务器通信。
后端存储 (Backend Store)
后端存储为每个运行持久化各种元数据,如运行 ID、开始和结束时间、参数、指标等。MLflow 支持两种类型的后端存储:基于文件系统(如本地文件)和基于数据库(如 PostgreSQL)。
此外,如果您正在与托管服务(如 Databricks 或 Azure Machine Learning)对接,您将与一个外部管理且无法直接访问的基于 REST 的后端存储进行交互。
制品存储 (Artifact Store)
制品存储为每个运行持久化(通常较大的)制品,如模型权重(例如 pickle 序列化的 scikit-learn 模型)、图像(例如 PNG)以及模型和数据文件(例如 Parquet 文件)。MLflow 默认将制品存储在本地文件 (mlruns) 中,但也支持不同的存储选项,如 Amazon S3 和 Azure Blob Storage。
对于作为 MLflow 制品记录的模型,您可以通过以下格式的模型 URI 引用该模型:models:/<model_id>,其中 'model_id' 是分配给记录模型的唯一标识符。这取代了旧的 runs:/<run_id>/<artifact_path> 格式,并提供了更直接的模型引用。
如果该模型已在 MLflow 模型注册表 (Model Registry) 中注册,您还可以通过以下格式的模型 URI 引用它:models:/<model-name>/<model-version>,详情请参阅 MLflow 模型注册表。
MLflow 追踪服务器 (Tracking Server)(可选)
MLflow 追踪服务器是一个独立的 HTTP 服务器,提供用于访问后端和/或制品存储的 REST API。追踪服务器还提供了配置服务内容、管理访问控制、版本控制等的灵活性。阅读 MLflow 追踪服务器文档了解更多详细信息。
常见设置
通过正确配置这些组件,您可以创建一个适合团队开发工作流的 MLflow 追踪环境。下表和图示展示了 MLflow 追踪环境的几种常见设置。
![]()
| 1. 本地主机(默认) | 2. 带本地数据库的本地追踪 | 3. 带 MLflow 追踪服务器的远程追踪 | |
|---|---|---|---|
| 场景 | 个人开发 | 个人开发 | 团队开发 |
| 用例 | 默认情况下,MLflow 将每次运行的元数据和制品记录到本地目录 mlruns。这是开始使用 MLflow Tracking 最简单的方法,无需设置任何外部服务器、数据库和存储。 | MLflow 客户端可以与兼容 SQLAlchemy 的数据库(例如 SQLite、PostgreSQL、MySQL)进行接口对接以用作后端。将元数据保存到数据库中,可以让您更清晰地管理实验数据,同时省去了设置服务器的工作。 | MLflow 追踪服务器可以配置制品 HTTP 代理,将制品请求通过追踪服务器转发,从而无需与底层的对象存储服务直接交互即可存储和检索制品。这对于团队开发场景特别有用,您可以将制品和实验元数据存储在带有适当访问控制的共享位置。 |
| 教程 | 快速入门 | 使用本地数据库追踪实验 | 使用 MLflow 追踪服务器进行远程实验追踪 |
其他使用 MLflow 追踪服务器的配置
MLflow 追踪服务器为其他特殊用例提供了可定制性。请遵循 使用 MLflow 追踪服务器进行远程实验追踪 来学习基本设置,并继续参考以下材料进行高级配置,以满足您的需求。
- 本地追踪服务器
- 仅制品模式
- 直接访问制品
在本地使用 MLflow 追踪服务器
当然,您可以在本地运行 MLflow 追踪服务器。虽然这相比直接使用本地文件或数据库并没有太大额外好处,但对于在本地测试团队开发工作流或在容器环境中运行机器学习代码可能很有用。
![]()
在仅制品模式下运行 MLflow 追踪服务器
MLflow 追踪服务器有一个 --artifacts-only 选项,允许服务器仅处理(代理)制品,而不允许处理元数据。这在大型组织或训练超大模型时特别有用。在这种场景下,您可能有大量的制品传输量,通过分流处理制品流量可以避免影响追踪功能。请阅读 选择性使用仅用于处理制品的追踪服务器实例 以获取有关如何使用此模式的更多详细信息。
![]()
禁用制品代理以允许直接访问制品
MLflow 追踪服务器默认同时服务制品和元数据。但在某些情况下,您可能希望允许直接访问远程制品存储,以避免代理的开销,同时保留元数据追踪功能。这可以通过在启动服务器时使用 --no-serve-artifacts 选项禁用制品代理来实现。有关如何设置此项,请参阅 在不代理制品访问的情况下使用追踪服务器。
![]()
常见问题 (FAQ)
我可以并行启动多个运行吗?
是的,MLflow 支持并行启动多个运行,例如多进程/线程。有关更多详细信息,请参阅 在同一个程序中启动多个运行。
我该如何整齐地组织众多的 MLflow 运行?
MLflow 提供了几种组织运行的方法:
- 将运行组织到实验中 - 实验是运行的逻辑容器。您可以使用 CLI、API 或 UI 创建实验。
- 创建子运行 - 您可以在单个父运行下创建子运行,将它们归为一组。例如,您可以为交叉验证实验中的每一折创建一个子运行。
- 为运行添加标签 - 您可以为每个运行关联任意标签,这允许您根据标签过滤和搜索运行。
我不运行追踪服务器可以直接访问远程存储吗?
可以。虽然在团队开发工作流中,使用 MLflow 追踪服务器作为制品访问代理是最佳实践,但如果您仅将其用于个人项目或测试,可能不需要这样做。您可以通过以下变通方法实现:
- 设置制品配置(如凭据和端点),就像为 MLflow 追踪服务器设置那样。详情请参阅 配置制品存储。
- 创建一个具有显式制品位置的实验,
experiment_name = "your_experiment_name"
mlflow.create_experiment(experiment_name, artifact_location="s3://your-bucket")
mlflow.set_experiment(experiment_name)
该实验下的运行将直接把制品记录到远程存储。
如何将 MLflow Tracking 与 模型注册表 (Model Registry) 集成?
要将模型注册表功能与 MLflow 追踪配合使用,您必须使用数据库后端存储(如 PostgreSQL),并使用相应模型口味 (model flavor) 的 log_model 方法记录模型。一旦模型被记录,您就可以通过 UI 或 API 在模型注册表中添加、修改、更新或删除该模型。有关如何为您的工作流正确配置后端存储,请参阅后端存储和常见设置。
如何包含关于运行的额外描述文本?
可以使用系统标签 mlflow.note.content 添加关于该运行的描述性说明。虽然其他系统标签是自动设置的,但此标签默认不会设置,用户可以覆盖它以包含关于运行的附加信息。内容将显示在运行页面的“笔记 (Notes)”部分。