MLflow 实验管理实践：从散乱 Notebook 到可复现的实验流水线
一、实验管理的混乱现状：为什么"能跑"不等于"能复现"
机器学习项目的早期阶段，实验通常在 Jupyter Notebook 中进行。模型参数写在代码注释里，训练数据路径硬编码在变量中，评估指标记录在电子表格里。当需要回溯"上周三那个 F1 达到 0.87 的实验用了什么参数"时，往往无从查起。这种混乱在单人多实验时已经令人头疼，在团队协作场景下更是灾难——不同成员用不同版本的训练脚本，数据集更新后旧实验结果无法复现。
MLflow 解决的核心问题是实验的可追溯性和可复现性。它提供四个模块：Tracking（实验记录）、Projects（项目打包）、Models（模型管理）、Registry（模型注册）。其中 Tracking 是使用频率最高的模块，它将每次实验的参数、指标和产物（Artifact）自动关联，形成完整的实验历史。
二、MLflow Tracking 架构：参数、指标与产物的关联模型
MLflow Tracking 的核心概念是 Run——一次完整的实验执行。每个 Run 记录三类信息：Parameters（超参数，如学习率、Batch Size）、Metrics（评估指标，如 Loss、Accuracy）、Artifacts（产物文件，如模型权重、混淆矩阵图）。多个 Run 组成一个 Experiment，Experiment 按 Name 组织。
flowchart TB
A[Experiment: 文本分类] –> B[Run 1: lr=0.001]
A –> C[Run 2: lr=0.0001]
A –> D[Run 3: lr=0.01]
B –> E[Params: lr=0.001, bs=32]
B –> F[Metrics: f1=0.87, loss=0.23]
B –> G[Artifacts: model.pt, cm.png]
C –> H[Params: lr=0.0001, bs=64]
C –> I[Metrics: f1=0.82, loss=0.31]
C –> J[Artifacts: model.pt, cm.png]
D –> K[Params: lr=0.01, bs=32]
D –> L[Metrics: f1=0.71, loss=0.52]
D –> M[Artifacts: model.pt, cm.png]
N[MLflow Server] –> A
N –> O[UI: 对比表格<br/>指标趋势图<br/>参数平行坐标图]
关键设计点：Metrics 支持时序记录（同一指标在不同 Step 的值），用于绘制训练曲线；Artifacts 存储在对象存储（S3、MinIO）中，不占用数据库空间；Run 之间可以通过 parent_run_id 形成层级关系，用于组织超参搜索中的子实验。
三、生产级代码实现：实验记录、超参搜索与模型注册
3.1 基础实验记录
import mlflow
import mlflow.pytorch
from mlflow.models.signature import infer_signature
def train_and_log(config, train_loader, val_loader):
"""训练模型并记录实验到 MLflow"""
# 为什么用 with 语法而非 start_run：with 块确保
# Run 在异常时也能正确关闭，避免僵尸 Run
with mlflow.start_run(
run_name=f"lr{config.lr}_bs{config.batch_size}"
) as run:
# 记录超参数
mlflow.log_params({
"learning_rate": config.lr,
"batch_size": config.batch_size,
"epochs": config.epochs,
"model_arch": config.model_name,
"seed": config.seed,
})
# 训练循环
model = build_model(config)
optimizer = torch.optim.AdamW(
model.parameters(), lr=config.lr)
best_f1 = 0.0
for epoch in range(config.epochs):
train_loss = train_one_epoch(
model, optimizer, train_loader)
val_metrics = evaluate(model, val_loader)
# 记录每个 Epoch 的指标（时序数据）
# 为什么按 Step 记录而非只记最终值：
# 训练曲线是判断过拟合和学习率调度的
# 关键依据，只看最终值无法诊断训练问题
mlflow.log_metrics({
"train_loss": train_loss,
"val_f1": val_metrics["f1"],
"val_precision": val_metrics["precision"],
"val_recall": val_metrics["recall"],
}, step=epoch)
if val_metrics["f1"] > best_f1:
best_f1 = val_metrics["f1"]
# 保存最优模型
torch.save(model.state_dict(),
"best_model.pt")
# 记录最终指标
mlflow.log_metric("best_val_f1", best_f1)
# 记录产物
mlflow.log_artifact("best_model.pt")
mlflow.log_artifact("confusion_matrix.png")
# 推断模型签名并记录
sample_input = next(iter(val_loader))[0][:1]
sample_output = model(sample_input)
signature = infer_signature(
sample_input.numpy(), sample_output.detach().numpy())
mlflow.pytorch.log_model(
model, "model", signature=signature)
return run.info.run_id
3.2 超参搜索与子 Run 管理
from mlflow.entities import ViewType
def hyperparameter_search(base_config, search_space,
train_loader, val_loader):
"""超参搜索：每个参数组合创建一个子 Run"""
# 创建父 Run，所有子 Run 挂在其下
# 为什么用父子 Run：超参搜索产生大量实验，
# 父子结构让 UI 中可以折叠/展开，
# 便于对比和筛选
with mlflow.start_run(
run_name="hparam_search"
) as parent_run:
mlflow.log_params({
"search_strategy": "grid",
"num_combinations": len(search_space),
})
best_run_id = None
best_f1 = 0.0
for params in search_space:
config = base_config.copy()
config.update(params)
# 创建子 Run
with mlflow.start_run(
run_name=f"lr{params['lr']}_bs{params['batch_size']}",
parent_id=parent_run.info.run_id
) as child_run:
run_id = train_and_log(
config, train_loader, val_loader)
# 查询子 Run 的最佳指标
run = mlflow.get_run(run_id)
f1 = run.data.metrics.get("best_val_f1", 0)
if f1 > best_f1:
best_f1 = f1
best_run_id = run_id
# 在父 Run 中记录搜索结果
mlflow.log_metric("search_best_f1", best_f1)
mlflow.log_param("best_run_id", best_run_id)
return best_run_id
3.3 模型注册与版本管理
def register_best_model(run_id, model_name,
stage="Staging"):
"""将最优模型注册到 Model Registry"""
# 注册模型版本
# 为什么需要 Model Registry：实验阶段产出的模型
# 需要经过评审才能上线，Registry 提供了
# Staging → Production 的阶段管理和审批流程
result = mlflow.register_model(
model_uri=f"runs:/{run_id}/model",
name=model_name
)
client = mlflow.tracking.MlflowClient()
# 将模型推进到 Staging
client.transition_model_version_stage(
name=model_name,
version=result.version,
stage=stage
)
# 添加模型描述
client.update_model_version(
name=model_name,
version=result.version,
description=f"基于 Run {run_id} 训练，"
f"F1 指标最优"
)
return result.version
3.4 实验对比与结果查询
def compare_experiments(experiment_name, metric="best_val_f1"):
"""查询实验结果并排序"""
client = mlflow.tracking.MlflowClient()
experiment = client.get_experiment_by_name(
experiment_name)
if experiment is None:
raise ValueError(
f"实验不存在: {experiment_name}")
# 搜索所有已完成的 Run
runs = client.search_runs(
experiment_ids=[experiment.experiment_id],
filter_string="",
run_view_type=ViewType.ACTIVE_ONLY,
order_by=[f"metrics.{metric} DESC"]
)
results = []
for run in runs[:10]: # 取 Top 10
results.append({
"run_id": run.info.run_id,
"params": run.data.params,
"metrics": run.data.metrics,
"start_time": run.info.start_time,
})
return results
四、MLflow 实践的架构权衡：存储后端、并发与数据治理
存储后端的选择：MLflow 支持 FileStore（本地文件系统）和 SQLStore（MySQL、PostgreSQL）作为后端存储。FileStore 适合单机实验，但无法支持多用户并发写入。SQLStore 是团队协作的必选项，但引入了数据库运维成本。Artifacts 存储同理——本地存储适合开发，S3/MinIO 适合生产。建议开发环境用 FileStore + 本地 Artifacts，生产环境用 SQLStore + S3。
并发写入的冲突问题：多人同时向同一 Experiment 写入 Run 时，可能出现 Run Name 冲突或指标覆盖。MLflow 的 Run ID 是全局唯一的，不会冲突，但 UI 展示时同名 Run 会造成混淆。建议在 Run Name 中加入用户标识和时间戳，或在 Experiment 层面按团队拆分。
数据版本与实验复现：MLflow 记录了参数和代码版本，但不自动管理数据版本。同一组参数在不同数据集上训练，结果可能完全不同。建议在参数中记录数据集的哈希值或 DVC 版本号，确保数据可追溯。
Model Registry 的审批流程局限：MLflow 的 Stage 转换（Staging → Production）缺乏内置的审批机制，任何有权限的人都可以推进。生产环境需要在外部实现审批逻辑（如通过 CI/CD Pipeline 触发），或在 MLflow 前面加一层 API 网关做权限控制。
五、总结
MLflow 实验管理的核心价值是将"散乱的 Notebook 实验"转化为"可追溯、可复现的实验流水线"。Tracking 模块解决参数和指标的关联记录，Projects 模块解决环境复现，Models 和 Registry 模块解决模型生命周期管理。落地时建议先从 Tracking 模块入手，在现有训练脚本中嵌入 MLflow 日志调用，再逐步引入超参搜索和模型注册。数据版本管理是实验复现的盲区，需要配合 DVC 或类似工具补齐。
