MLFlow：模型版本控制与实验管理

1.1 为什么需要 MLFlow：先说问题

MLFlow 四大核心组件架构——Tracking/Projects/Models/Registry 的关系与功能

想象这个场景：你花了两周时间跑了三十个实验，发现第 17 个实验的结果最好。三个月后，你的同事想复现这个结果——怎么办？

如果没有专门的工具，答案通常是"找不到了"。参数不知道，数据版本不确定，代码可能改过了，结果无法复现。

这是 AI 开发和传统软件开发最大的不同之一：代码相同，但数据+参数+随机种子不同，模型就不同。Git 只能管代码，管不了这些。MLFlow 就是解决这个问题的。

MLFlow 解决的三个具体问题：

实验记录混乱：每次训练的参数是什么？指标是多少？用了哪版数据？MLFlow 的 Tracking 模块把这些都记录下来，形成可查询的实验历史。

模型版本无法追踪：哪个模型在生产？哪个在测试？上个月的最好模型在哪里？MLFlow Registry 管理模型的生命周期：Staging（测试中）→ Production（生产）→ Archived（归档）。

实验无法对比：这次改了学习率和 batch size，效果好还是不好？MLFlow UI 提供实验对比界面，参数和指标一目了然。

1.2 安装和启动

bash

pip install mlflow

# 启动 MLFlow UI（会在当前目录创建 mlruns 文件夹存储数据）
mlflow ui

# 访问 http://localhost:5000 查看实验结果

1.3 基础用法：记录一次实验

MLFlow 的使用模式很简单：用 mlflow.start_run() 开启一次实验运行，在里面记录参数、指标和模型，结束后在 UI 里查看。

python

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score

# 加载数据
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
    iris.data, iris.target, test_size=0.2, random_state=42
)

# 设置实验名称（相当于项目名，多次运行归属到同一实验下）
mlflow.set_experiment("iris-classification")

# 开始一次实验运行
with mlflow.start_run(run_name="random-forest-v1"):
    # 1. 记录参数（超参数、配置等）
    params = {"n_estimators": 100, "max_depth": 5, "random_state": 42}
    mlflow.log_params(params)

    # 2. 训练模型
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)

    # 3. 记录指标（评估结果）
    y_pred = model.predict(X_test)
    metrics = {
        "accuracy": accuracy_score(y_test, y_pred),
        "f1_macro": f1_score(y_test, y_pred, average='macro')
    }
    mlflow.log_metrics(metrics)

    # 4. 保存模型（保存到 MLFlow 存储，可以后续加载）
    mlflow.sklearn.log_model(model, "random-forest-model")

    print(f"准确率: {metrics['accuracy']:.4f}")
    print(f"F1分数: {metrics['f1_macro']:.4f}")
    print(f"实验运行已记录，在 http://localhost:5000 查看")

运行后，打开 http://localhost:5000 就能看到这次实验的所有信息：参数、指标、保存的模型文件，还有运行时间。

1.4 对比多次实验：找到最优超参数

MLFlow 最有价值的场景是超参数搜索的记录和对比。你改了参数跑了几十次，通过 UI 能直接看出哪组参数最好，不用手动整理 Excel。

python

# 对比不同超参数的效果
experiments_config = [
    {"n_estimators": 50, "max_depth": 3},
    {"n_estimators": 100, "max_depth": 5},
    {"n_estimators": 200, "max_depth": 10},
    {"n_estimators": 200, "max_depth": 15},
]

mlflow.set_experiment("iris-hyperparameter-search")

for i, params in enumerate(experiments_config):
    with mlflow.start_run(run_name=f"experiment-{i+1}"):
        mlflow.log_params(params)

        model = RandomForestClassifier(**params, random_state=42)
        model.fit(X_train, y_train)

        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        mlflow.log_metric("accuracy", accuracy)

        mlflow.sklearn.log_model(model, "model")
        print(f"实验{i+1}: 参数={params}, 准确率={accuracy:.4f}")

# 用代码找最佳实验（也可以在 UI 里点击对比）
runs = mlflow.search_runs(experiment_names=["iris-hyperparameter-search"])
best_run = runs.loc[runs["metrics.accuracy"].idxmax()]
print(f"\n最佳实验:")
print(f"  准确率: {best_run['metrics.accuracy']:.4f}")
print(f"  参数: n_estimators={best_run['params.n_estimators']}, max_depth={best_run['params.max_depth']}")

1.5 模型注册和版本管理

为什么需要模型注册表？

想象你有三个版本的模型：v1 在生产，v2 正在测试，v3 刚训练完。没有版本管理工具，你怎么确定"现在生产在用哪个"？怎么回滚到 v1？怎么知道各个版本的性能数据？

MLFlow Model Registry 解决了这个问题，它的核心是给模型版本定义生命周期状态：

code

None（刚注册）→ Staging（测试中）→ Production（生产）→ Archived（已归档）

python

import mlflow
from mlflow.tracking import MlflowClient

client = MlflowClient()

# 把最优实验的模型注册到 Model Registry
# 先找到 run_id（从上面的搜索结果里取，或者从 UI 里复制）
best_run_id = best_run["run_id"]
model_uri = f"runs:/{best_run_id}/model"

# 注册模型（首次注册会创建模型，后续注册会创建新版本）
model_details = mlflow.register_model(model_uri, "IrisClassifier")
print(f"模型已注册: 版本 {model_details.version}")

# 将版本 1 推进到 Staging（测试环境）
client.transition_model_version_stage(
    name="IrisClassifier",
    version=1,
    stage="Staging"
)
print("模型已进入 Staging 阶段")

# 测试验证通过后，推进到 Production
client.transition_model_version_stage(
    name="IrisClassifier",
    version=1,
    stage="Production"
)
print("模型已上线 Production")

# 加载生产版本的模型（不需要知道具体版本号）
production_model = mlflow.sklearn.load_model("models:/IrisClassifier/Production")
print(f"加载生产模型，预测示例: {production_model.predict(X_test[:3])}")

Model Registry 的使用规范：

python

# 查看模型的所有版本
versions = client.search_model_versions("name='IrisClassifier'")
for v in versions:
    print(f"版本 {v.version}: {v.current_stage} - {v.description}")

# 为模型版本添加描述（记录该版本的变更内容）
client.update_model_version(
    name="IrisClassifier",
    version=2,
    description="使用更大的 n_estimators=200，准确率提升 1.2%"
)

# 添加标签（用于过滤和搜索）
client.set_model_version_tag(
    name="IrisClassifier",
    version=2,
    key="validated_by",
    value="alice"
)

1.6 记录 LLM 微调实验

MLFlow 对 LLM 微调场景同样适用。关键是把微调过程中的所有关键信息记录下来，确保实验可复现。

python

import mlflow
import mlflow.pytorch
import torch

mlflow.set_experiment("llm-finetuning")

with mlflow.start_run(run_name="qwen-lora-v1"):
    # 记录微调的完整配置（这是复现的关键）
    mlflow.log_params({
        "base_model": "Qwen/Qwen2.5-7B-Instruct",
        "base_model_hash": "abc123",     # 模型文件的 commit hash，确保用了同一个版本
        "lora_rank": 16,
        "lora_alpha": 32,
        "lora_target_modules": "q_proj,k_proj,v_proj,o_proj",
        "learning_rate": 2e-4,
        "lr_scheduler": "cosine",
        "num_epochs": 3,
        "batch_size": 4,
        "gradient_accumulation_steps": 4,  # 等效 batch_size = 4*4 = 16
        "max_seq_length": 1024,
        "training_data_path": "training_data_final.json",
        "training_data_size": 5000,
        "training_data_hash": "def456",    # 数据文件的哈希，确保数据版本一致
    })

    # 训练过程中，每个 epoch 记录 loss（这是最重要的训练过程指标）
    for epoch in range(3):
        # 实际训练代码在这里
        # trainer.train() ...

        # 记录每轮的指标（step 参数让 MLFlow 知道时间顺序）
        train_loss = 0.5 - epoch * 0.12   # 示例数据
        eval_loss = 0.6 - epoch * 0.10    # 示例数据

        mlflow.log_metrics({
            "train_loss": train_loss,
            "eval_loss": eval_loss
        }, step=epoch)

    # 训练结束后，记录最终评估结果
    mlflow.log_metrics({
        "final_eval_loss": 0.34,
        "bleu_score": 0.72,
        "human_eval_score": 4.2,   # 人工评估分数（1-5）
        "gpt4_eval_score": 7.8     # GPT-4 自动评分（1-10）
    })

    # 记录 LoRA 权重（作为 artifact 保存）
    mlflow.log_artifact("./lora-weights")

    # 记录重要的配置文件（便于复现）
    mlflow.log_artifact("./lora-config.json")

    print("LLM 微调实验已记录")

1.7 实验追踪的最佳实践

使用 MLFlow 时，这几条原则能让它真正发挥价值：

记录一切，不只是你认为重要的。数据文件哈希、代码 commit hash、环境依赖版本——这些看起来无关紧要，但三个月后想复现时你会感激自己当初记录了。

给每次运行起有意义的名字。run_name="qwen-lora-v1-lr2e4-r16" 比自动生成的随机字符串有用多了，不用打开详情就能大致知道这次实验做了什么。

写描述，不要依赖记忆。client.update_model_version(description="使用更大数据集，修复了格式不一致问题") 这类描述在几个月后还能告诉你这个版本的来龙去脉。

不要在 Production 直接更新。应该先让新版本进 Staging，跑一段时间验证没问题，再升级到 Production。可以快速回滚到上一个 Production 版本。

1.8 小结

MLFlow 解决的是 AI 开发中最常见但最难管理的工程问题：

实验记录：每次训练的参数、指标、代码版本，一次都不丢
模型版本：明确知道哪个版本在生产，可以快速回滚
可复现性：能精确复现任何一次实验（只要你记录了数据版本）

"上次效果更好的模型参数是什么？"——这个问题从此有了答案。