MLflow 跟踪

MLflow Tracking 是一个用于在运行你的机器学习代码时记录参数、代码版本、指标和输出文件的 API 和 UI，并可在之后可视化这些结果。MLflow Tracking 提供 Python、REST、R 和 Java APIs。

快速开始​

如果您以前没有使用过 MLflow Tracking，我们强烈建议先查看下面的快速入门教程。

概念​

运行​

MLflow Tracking 是围绕 运行 概念组织的，这些运行是某段数据科学代码的执行实例，例如，单次 python train.py 执行。每个运行会记录元数据（有关运行的各种信息，例如指标、参数、开始和结束时间）和工件（运行产生的输出文件，例如模型权重、图像等）。

模型​

模型表示在运行期间产生的已训练的机器学习工件。已记录的模型包含与运行类似的、属于它们自身的元数据和工件。

实验​

实验将针对特定任务把运行和模型分组在一起。你可以使用 CLI、API 或 UI 创建实验。 MLflow 的 API 和 UI 也允许你创建并搜索实验。有关如何将你的运行组织到实验中的更多详细信息，请参见 Organizing Runs into Experiments。

跟踪运行​

MLflow Tracking APIs 提供了一组函数来跟踪你的运行。 例如，你可以调用 mlflow.start_run() 来开始一个新运行， 然后调用 Logging Functions，比如 mlflow.log_param() 和 mlflow.log_metric() 来分别记录参数和指标。 请访问 Tracking API documentation 以获取有关使用这些 APIs 的更多详细信息。

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()。Auto-logging 支持流行的库，例如 Scikit-learn、XGBoost、PyTorch、 Keras、Spark 等。查看 Automatic Logging Documentation 以了解受支持的库以及如何对每个库使用自动记录 API。

import mlflow

mlflow.autolog()

# Your training code...

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

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:/ URI 格式以便立即访问模型

有关全面的示例和高级搜索模式，请参阅 Search Logged Models Guide。

以编程方式查询运行​

你也可以使用 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)
# {'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"

这种新方法提供了若干优势：

• 直接模型引用: 无需知道 run ID 和 artifact path
• 更好的模型生命周期管理：每个模型检查点都有自己的唯一标识符
• 改进的模型比较：轻松比较同一次运行内的不同检查点
• 增强的可追溯性: 模型、指标和数据集之间的清晰关联

跟踪数据集​

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

探索运行、模型和结果​

跟踪界面​

Tracking UI 让您以可视化方式浏览您的实验、运行和模型，如本页顶部所示。

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

这些功能对于模型也可用，如下所示。

如果你将运行记录保存到本地 mlruns 目录，请在其上级目录中运行以下命令，然后在浏览器中访问 http://127.0.0.1:5000。

mlflow ui --port 5000

或者，MLflow Tracking Server 提供相同的 UI 并支持运行产物的远程存储。在这种情况下，您可以在任何能够连接到您的跟踪服务器的机器上通过 http://:5000 查看该 UI。

设置 MLflow 跟踪环境​

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

组件​

MLflow Tracking APIs​

你可以在你的机器学习代码中调用 MLflow Tracking APIs 来记录运行并在必要时与 MLflow Tracking Server 通信。

后端存储​

后端存储会为每个 Run 持久化各种元数据，例如运行 ID、开始和结束时间、参数、指标等。MLflow 支持两种后端存储类型：基于文件系统的（如本地文件）和 基于数据库的（如 PostgreSQL）。

此外，如果您正在与托管服务（例如 Databricks 或 Azure Machine Learning）交互，您将与一个基于 REST 的后端存储进行交互，该存储是外部管理的，无法直接访问。

Artifact Store​

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

对于作为 MLflow 工件记录的模型，您可以通过以下格式的模型 URI 引用该模型：models:/， 其中 'model_id' 是分配给已记录模型的唯一标识符。该格式取代了较旧的 runs:// 格式， 并提供了更直接的模型引用方式。

如果模型已注册在 MLflow Model Registry，也可以通过格式为的模型 URI 来引用该模型：models://，详情请参见 MLflow Model Registry。

MLflow Tracking Server (可选)​

MLflow Tracking Server 是一个独立的 HTTP 服务器，提供用于访问后端和/或 artifact 存储的 REST APIs。该服务器还提供灵活性，可配置要提供的数据、管理访问控制、版本控制等。阅读 MLflow Tracking Server documentation 了解更多详情。

常见设置​

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

与 MLflow Tracking Server​ 的其他配置

MLflow Tracking Server 为其他特殊用例提供定制能力。请参阅 Remote Experiment Tracking with MLflow Tracking Server 来学习基本设置，并继续查阅以下材料以获取满足您需求的高级配置。

常见问题​

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

是的，MLflow 支持并行启动多个运行，例如多进程/多线程。详情请参见 Launching Multiple Runs in One Program。

我如何将大量 MLflow Runs 整理得更整洁？​

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

• Organize runs into experiments - 实验是用于存放你的运行的逻辑容器。你可以使用 CLI、API 或 UI 创建一个实验。
• Create child runs - 您可以在单个父运行下创建子运行以将它们分组。例如，您可以为交叉验证实验中的每个折创建一个子运行。
• Add tags to runs - 你可以为每个运行关联任意标签，从而根据标签筛选和搜索运行。

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

是的，尽管在团队开发工作流中，将 MLflow Tracking Server 作为工件访问的代理是最佳实践，但如果您用于个人项目或测试，可能不需要这样做。您可以通过下面的变通方法实现：

1. 设置 artifacts 的配置，例如凭证和端点，就像为 MLflow Tracking Server 所做的那样。请参见 configure artifact storage 以获取更多详情。
2. 创建具有显式工件位置的实验，

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）并使用相应模型风格的 log_model 方法记录模型。一旦模型被记录，您可以通过 UI 或 API 在模型注册中添加、修改、更新或删除该模型。参见 Backend Stores 和 Common Setups 以了解如何为您的工作流正确配置后端存储。

如何包含关于运行的额外描述文本？​

系统标签 mlflow.note.content 可用于为此运行添加描述性注释。虽然其他 system tags 会自动设置， 但该标签 默认不会设置，用户可以覆盖它以包含有关该运行的附加信息。该内容将显示在运行页面的 Notes 部分。