📚 课程目录第4章使用MLflow跟踪实验

4. 使用MLflow跟踪实验

在 Databricks 上做机器学习，最容易被低估的环节不是建模，而是“把每一次尝试都记清楚”。MLflow Tracking 正是为此设计的记录与查询组件：它把一次训练过程抽象为一次 Run（运行），把 Run 归入某个 Experiment（实验），并围绕 Run 记录参数、指标、工件、标签和源码环境信息。理解这套对象模型，是后续模型注册、部署与回滚的基础。

一、四个核心对象：Experiment、Run、Param、Metric、Artifact

MLflow Tracking 的数据模型可以概括为“两层容器 + 三类记录”。Experiment 是最外层容器，通常对应一个业务问题或一个模型方向，例如 /Users/mei@corp.com/churn-prediction。它下面挂载多个 Run，每个 Run 代表一次独立的训练或评估执行。Run 内部再记录三类内容：Param（输入侧的配置，如 max_depth=6、learning_rate=0.05）、Metric（输出侧的数值结果，如 auc=0.873、log_loss=0.412）、Artifact（文件类产物，如模型文件、混淆矩阵图、特征重要性 CSV、conda 环境文件）。

需要特别区分 Param 与 Metric 的语义：Param 是“你设定的”，一次 Run 内同名参数只保留一个值；Metric 是“训练算出来的”，同名指标可以按 step 记录多次，形成一条随训练推进的曲线。这个区别直接决定了你在 UI 里能否画出学习曲线，也决定了后续对比实验时哪些字段可排序、可筛选。

二、在 Databricks 中创建实验并组织运行

Databricks 的 MLflow 集成有两种常见入口：Notebook 内的 mlflow Python API，以及左侧栏的 Experiments 页面。推荐做法是“先建实验、再跑运行”，而不是依赖默认实验。默认实验通常以 Notebook 路径命名，容易把不同业务方向的 Run 混在一起，后期筛选成本很高。

import mlflow
import mlflow.sklearn
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import roc_auc_score, log_loss

# 1) 指定或创建实验（幂等：已存在则直接返回其 ID）
exp_name = "/Users/mei@corp.com/churn-prediction"
mlflow.set_experiment(exp_name)

# 2) 开启一次 Run，并打上便于检索的标签
with mlflow.start_run(run_name="gbt-lr005-depth6") as run:
    mlflow.set_tags({
        "team": "risk-modeling",
        "data_version": "2024-06-01",
        "stage": "baseline"
    })

    params = {"max_depth": 6, "learning_rate": 0.05, "n_estimators": 300}
    mlflow.log_params(params)

    model = GradientBoostingClassifier(**params)
    model.fit(X_train, y_train)

    proba = model.predict_proba(X_valid)[:, 1]
    mlflow.log_metric("auc", roc_auc_score(y_valid, proba))
    mlflow.log_metric("log_loss", log_loss(y_valid, proba))

    mlflow.sklearn.log_model(model, artifact_path="model")
    mlflow.log_artifact("reports/confusion_matrix.png")

    print("run_id =", run.info.run_id)

这段代码体现了几个关键约定。set_experiment 负责把后续 Run 归入指定实验；start_run 作为上下文管理器，退出时自动把 Run 状态置为 FINISHED，若中途抛异常则置为 FAILED，这对排查“哪些尝试没跑完”非常有用。log_params 一次性写入多个参数，log_metric 逐条写入指标，log_model 与 log_artifact 分别写入模型和普通文件。运行结束后，run.info.run_id 就是这次尝试的全局唯一标识，后续加载模型、注册模型、做对比都靠它。

三、参数、指标、工件的记录要点

参数记录要遵循“可复现”原则：凡是影响训练结果的输入都应记录，包括超参数、随机种子、数据版本标识、特征集名称。不要把大段文本塞进参数，参数值应短小、可比较。若某个配置项很长（如完整特征列表），更适合作为标签或工件文件保存。

指标记录要区分“单值指标”和“序列指标”。像最终 AUC 只需记录一次；像每轮迭代的 loss 则应带 step 参数多次记录，例如 mlflow.log_metric("train_loss", v, step=epoch)。这样在 Experiments UI 中可以直接看到曲线，判断是否过拟合或收敛。指标名建议统一命名规范，如 valid_auc、train_auc，避免同一含义出现多种写法导致无法聚合对比。

工件记录要控制体积与结构。模型文件、评估图表、特征重要性、环境文件都适合作为工件；但不要把原始训练数据整体上传，那会让实验存储迅速膨胀。Databricks 上工件默认存放在 DBFS 的实验目录下，可通过 mlflow.log_artifact(local_path, artifact_path) 指定子目录，例如统一放到 plots/ 或 reports/，便于浏览。

在 Databricks Notebook 中，mlflow.start_run() 不传参数时会自动复用当前 Notebook 的活跃 Run；若想强制新建，可加 nested=True 或先 mlflow.end_run()。超参数调优时常用嵌套 Run 表达“父调优任务 + 子单次试验”的层级关系。

四、用标签、搜索与对比把实验“组织”起来

实验数量一多，能否快速找到“上周风险组用 6 月数据跑出的 AUC 最高的那次运行”就成了效率瓶颈。MLflow 提供了基于标签和指标的搜索语法，Databricks UI 的搜索框与 mlflow.search_runs API 都支持。例如按标签筛选：

runs = mlflow.search_runs(
    experiment_names=["/Users/mei@corp.com/churn-prediction"],
    filter_string="tags.team = 'risk-modeling' and metrics.auc > 0.85",
    order_by=["metrics.auc DESC"],
    max_results=10
)
print(runs[["run_id", "params.max_depth", "metrics.auc"]])

这段查询返回一个 pandas DataFrame，列名以 params.、metrics.、tags. 为前缀，可以直接排序、导出、画图。实践中建议固定三类标签：team（归属）、data_version（数据快照）、stage（baseline / tuning / final）。这样无论实验数量增长到几百还是几千，都能通过标签组合快速收敛候选集，再在 UI 中勾选多个 Run 做并排对比，查看参数差异、指标曲线和工件。

常见误区：把实验当成“文件夹”随意嵌套，或把不同业务方向的 Run 全塞进默认实验。实验路径应稳定、语义清晰；一旦实验被大量下游引用（如模型注册、自动化任务），重命名或迁移会带来额外维护成本。

五、一次完整跟踪的检查清单

为了让每次 Run 都“可复现、可比较、可追溯”，可以在训练脚本收尾前对照以下清单：实验路径是否明确且非默认；Run 是否有可读的 run_name；参数是否覆盖全部影响结果的输入；指标是否包含验证集与训练集的关键度量；模型与评估图表是否已作为工件保存；标签是否包含团队、数据版本、阶段；Run 状态是否为 FINISHED。满足这些条件后，这次运行才真正具备进入模型注册与部署流程的资格。

MLflow Tracking 的价值不在于“记录了多少”，而在于“需要时能否快速找到并复现”。把实验、运行、参数、指标、工件这五个对象用对、用稳，后续的模型版本管理、生产部署与回滚才有可靠的事实依据。