3. MLflow 基础入门：Tracking 核心 API 实战

说实话，我第一次接触 MLflow 的时候，心里想的是：「又一个实验管理工具？」。但真正用起来之后，我发现它确实不一样。它不强迫你改变代码结构，也不要求你迁移到某个特定平台。你只需要在训练脚本里加几行代码，就能把参数、指标、模型文件全部记录下来。

这一章，我们就来亲手操作 MLflow Tracking 的四个核心 API。我会带着你，从最基础的参数记录开始，一步步搭建起一个可追溯的实验记录体系。

3.1 MLflow Tracking 是什么？

MLflow Tracking 是 MLflow 的组件之一。它的核心功能就一句话：记录和查询机器学习实验的元数据。

这些元数据包括：

• 参数（Parameters）：你训练模型时设置的超参数，比如学习率、批量大小、层数。
• 指标（Metrics）：模型在训练和验证过程中的表现，比如准确率、损失值、F1 分数。
• 标签（Tags）：用来标记实验的元信息，比如实验名称、运行环境、Git 提交 ID。
• 工件（Artifacts）：模型文件、图表、日志文件、数据集样本等任何文件。

我个人习惯把 MLflow Tracking 想象成一个「实验日记本」。每次你运行一次训练，它就自动帮你记下：你用了什么参数、跑出了什么结果、产出了哪些文件。以后你想复盘，翻开这本日记本，一目了然。

3.2 环境准备与第一个 Run

在开始之前，你需要安装 MLflow。很简单，一行命令搞定：

pip install mlflow

安装完成后，我们来创建第一个实验记录。我会用一个简单的线性回归训练作为例子。

3.3 记录参数：log_param

参数是实验的「输入」。你想想看，如果连学习率是多少都忘了，那实验结果还有什么参考价值？

log_param 的用法非常直接：

import mlflow

# 创建一个新的 Run
with mlflow.start_run():
    # 记录参数
    mlflow.log_param("learning_rate", 0.01)
    mlflow.log_param("batch_size", 32)
    mlflow.log_param("epochs", 50)
    mlflow.log_param("optimizer", "adam")

这里有几个要点：

• 参数名必须是字符串，值可以是字符串、数字或布尔值。
• 同一个 Run 中，参数名不能重复。如果你重复记录同名参数，后面的值会覆盖前面的。
• 参数一旦记录，就不能修改。这是设计上的考虑——实验参数应该是不可变的。

params = {
    "learning_rate": 0.01,
    "batch_size": 32,
    "epochs": 50,
    "optimizer": "adam"
}
for key, value in params.items():
    mlflow.log_param(key, value)

3.4 记录指标：log_metric

指标是实验的「输出」。它告诉你模型训练得怎么样。与参数不同，指标是可以随时间变化的——比如每个 epoch 的损失值。

log_metric 的基本用法：

with mlflow.start_run():
    # 模拟训练过程
    for epoch in range(10):
        loss = 1.0 / (epoch + 1)  # 模拟损失下降
        accuracy = 0.5 + epoch * 0.05  # 模拟准确率上升
        
        # 记录每个 epoch 的指标
        mlflow.log_metric("loss", loss, step=epoch)
        mlflow.log_metric("accuracy", accuracy, step=epoch)

注意这里多了一个 step 参数。它表示当前指标是在第几步记录的。MLflow 会根据 step 值绘制指标变化曲线。如果你不传 step，默认从 0 开始递增。

我在项目中遇到过一个问题：有时候训练到一半中断了，但之前记录的指标还在。这其实是 MLflow 的一个优点——指标是增量记录的，不会因为训练中断而丢失之前的数据。

3.5 记录标签：set_tag

标签是用来给实验做「标记」的。它不像参数那样影响模型训练，也不像指标那样反映模型表现。它的作用是帮你快速筛选和定位实验。

常见的标签用法：

with mlflow.start_run():
    mlflow.set_tag("experiment_name", "lr_search")
    mlflow.set_tag("model_type", "linear_regression")
    mlflow.set_tag("dataset", "housing_prices")
    mlflow.set_tag("git_commit", "a1b2c3d4e5")
    mlflow.set_tag("run_owner", "zhang_san")

我个人习惯把标签分为两类：

• 业务标签：实验名称、数据集版本、模型类型等。
• 环境标签：Git 提交 ID、运行机器、Python 版本等。

你想想看，当你一个月后面对几百个 Run 时，如果没有标签，你根本不知道哪个 Run 对应哪个实验。有了标签，你可以用 MLflow UI 直接按标签筛选，效率高得多。

3.6 记录工件：log_artifact

工件是实验的「产出」。模型文件、训练日志、损失曲线图、特征重要性图……只要是文件，都可以作为工件记录下来。

log_artifact 的用法：

import mlflow
import matplotlib.pyplot as plt

# 先训练模型并保存
model = train_model()
model.save("model.pkl")

# 绘制损失曲线并保存
plt.plot(loss_history)
plt.savefig("loss_curve.png")

# 记录工件
with mlflow.start_run():
    mlflow.log_artifact("model.pkl")
    mlflow.log_artifact("loss_curve.png")
    mlflow.log_artifact("config.yaml")

这里有个细节：log_artifact 默认会把文件复制到 MLflow 的存储目录中。你不用担心原始文件被删除后记录丢失。

如果你想记录整个目录，可以用 log_artifacts（注意是复数）：

mlflow.log_artifacts("outputs/")  # 记录 outputs 目录下的所有文件

3.7 完整示例：一个可追溯的训练脚本

好了，我们把上面学到的 API 组合起来，写一个完整的训练脚本。这个脚本运行一次，就会生成一个完整的实验记录。

import mlflow
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error

# 生成模拟数据
X = np.random.rand(100, 1) * 10
y = 2 * X + 1 + np.random.randn(100, 1) * 0.5

# 定义超参数
params = {
    "fit_intercept": True,
    "normalize": False,
    "copy_X": True
}

# 开始实验记录
with mlflow.start_run():
    # 1. 记录参数
    for key, value in params.items():
        mlflow.log_param(key, value)
    
    # 2. 训练模型
    model = LinearRegression(**params)
    model.fit(X, y)
    
    # 3. 计算指标
    predictions = model.predict(X)
    mse = mean_squared_error(y, predictions)
    rmse = np.sqrt(mse)
    
    mlflow.log_metric("mse", mse)
    mlflow.log_metric("rmse", rmse)
    
    # 4. 记录标签
    mlflow.set_tag("model_type", "linear_regression")
    mlflow.set_tag("dataset", "synthetic_data")
    mlflow.set_tag("task", "regression")
    
    # 5. 保存并记录模型
    import joblib
    joblib.dump(model, "linear_model.pkl")
    mlflow.log_artifact("linear_model.pkl")
    
    print(f"Run ID: {mlflow.active_run().info.run_id}")
    print(f"MSE: {mse:.4f}, RMSE: {rmse:.4f}")

运行这段代码后，你会在终端看到 Run ID。记下它，后面我们会用到。

3.8 查看记录：MLflow UI

记录完实验后，怎么查看呢？MLflow 提供了一个 Web UI。在终端运行：

mlflow ui

然后在浏览器打开 http://localhost:5000。你会看到一个简洁的界面，里面列出了所有 Run。点击某个 Run，就能看到它的参数、指标、标签和工件。

我个人觉得，MLflow UI 最实用的功能是「对比」。你可以勾选多个 Run，然后并排查看它们的参数和指标。哪个参数组合效果最好，一目了然。

3.9 避坑指南

最后，分享几个我踩过的坑：

• 不要在一个脚本里多次调用 mlflow.start_run()。如果你需要记录多个独立的实验，用 with mlflow.start_run() 块来隔离。嵌套使用会导致混乱。
• 指标名称不要用中文。虽然 MLflow 支持 Unicode，但 UI 的排序和筛选功能对中文支持不太好。我习惯用英文命名指标。
• 工件路径不要包含特殊字符。空格、中文、特殊符号可能导致某些系统上文件读取失败。我一般只用字母、数字和下划线。
• 及时清理无用的 Run。实验多了之后，存储会膨胀。我每个月会清理一次，只保留有意义的实验记录。