MLOps 实战：模型版本管理与实验追踪

# 没有实验追踪的典型混乱场景
models = {
    "final_v1.pkl": "不知道用的什么参数",
    "final_v2.pkl": "这个好像更好？",
    "best_model.pkl": "等等，哪个是best？",
    "model_20260301.pkl": "这是哪天训练的？",
    "model_final_final.pkl": "这个真的是final吗...",
}
# 三个月后：完全无法复现任何结果

# 使用 MLflow 追踪实验
import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

mlflow.set_experiment("customer-churn-prediction")

with mlflow.start_run():
    # 自动记录参数
    mlflow.log_param("n_estimators", 100)
    mlflow.log_param("max_depth", 10)
    mlflow.log_param("random_state", 42)
    
    # 训练模型
    model = RandomForestClassifier(
        n_estimators=100, max_depth=10, random_state=42
    )
    model.fit(X_train, y_train)
    
    # 记录指标
    accuracy = accuracy_score(y_test, model.predict(X_test))
    mlflow.log_metric("accuracy", accuracy)
    
    # 保存模型
    mlflow.sklearn.log_model(model, "model")
    
    print(f"Accuracy: {accuracy:.4f}")
    print(f"Run ID: {mlflow.active_run().info.run_id}")

# MLflow 实验层级结构
import mlflow

# 创建实验
mlflow.set_experiment("recommendation-system")

# 查看实验信息
experiment = mlflow.get_experiment_by_name("recommendation-system")
print(f"Experiment ID: {experiment.experiment_id}")
print(f"Artifact Location: {experiment.artifact_location}")

# 搜索历史 runs
runs = mlflow.search_runs(
    experiment_ids=[experiment.experiment_id],
    filter_string="metrics.accuracy > 0.9",
    order_by=["metrics.accuracy DESC"]
)
print(f"找到 {len(runs)} 个高精度实验")

# 启动 MLflow UI 查看实验
mlflow ui --host 0.0.0.0 --port 5000

# 或指定后端存储
mlflow server \
    --backend-store-uri sqlite:///mlflow.db \
    --default-artifact-root ./artifacts \
    --host 0.0.0.0 \
    --port 5000

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

import wandb
import torch
import torch.nn as nn

# 初始化实验
wandb.init(
    project="image-classification",
    config={
        "learning_rate": 0.001,
        "epochs": 50,
        "batch_size": 32,
        "model": "resnet18",
    }
)

# 自动记录配置
config = wandb.config

# 训练循环中记录指标
for epoch in range(config.epochs):
    train_loss = train_one_epoch(model, train_loader)
    val_acc = evaluate(model, val_loader)
    
    wandb.log({
        "train_loss": train_loss,
        "val_accuracy": val_acc,
        "epoch": epoch,
    })

# WandB 超参数扫描（Sweep）
sweep_config = {
    "method": "bayes",  # 贝叶斯优化
    "metric": {"name": "val_accuracy", "goal": "maximize"},
    "parameters": {
        "learning_rate": {"min": 1e-5, "max": 1e-2, "distribution": "log_uniform"},
        "batch_size": {"values": [16, 32, 64, 128]},
        "hidden_dim": {"values": [64, 128, 256, 512]},
    }
}

# 创建 sweep
sweep_id = wandb.sweep(sweep_config, project="hyperparameter-search")

# 运行 agent
wandb.agent(sweep_id, function=train, count=50)
# 自动运行 50 次实验，寻找最优超参数

# 初始化 DVC
dvc init

# 跟踪数据集
dvc add data/train.csv
dvc add data/test.csv

# 提交到 Git（只提交元文件）
git add data/train.csv.dvc data/test.csv.dvc .dvc
git commit -m "Add dataset v1.0"

# 推送数据到远程存储
dvc remote add -d myremote s3://my-bucket/dvc-store
dvc push

# 在不同数据集版本之间切换
git checkout experiment-v2  # 切换到另一个分支
dvc checkout               # 切换对应的数据

# 比较数据集变化
dvc diff HEAD~1

# 复现特定实验
git checkout abc1234  # 切换到特定 commit
dvc checkout          # 恢复对应的数据集
python train.py       # 运行训练

# MLflow Projects 定义实验流程
# MLproject 文件
name: churn-prediction
conda_env: conda.yaml

entry_points:
  main:
    parameters:
      n_estimators: {type: int, default: 100}
      max_depth: {type: int, default: 10}
      data_version: {type: string, default: "v1.0"}
    command: "python train.py {n_estimators} {max_depth} {data_version}"

# 运行项目
mlflow run . -P n_estimators=200 -P max_depth=15 -P data_version=v2.0

# 自动化实验调度
import mlflow
import itertools
import pandas as pd

param_grid = {
    "n_estimators": [50, 100, 200],
    "max_depth": [5, 10, 15, None],
    "min_samples_split": [2, 5, 10],
}

mlflow.set_experiment("hyperparameter-grid-search")
results = []

for params in itertools.product(*param_grid.values()):
    param_dict = dict(zip(param_grid.keys(), params))
    
    with mlflow.start_run():
        for k, v in param_dict.items():
            mlflow.log_param(k, v)
        
        model = RandomForestClassifier(param_dict)
        model.fit(X_train, y_train)
        accuracy = model.score(X_test, y_test)
        mlflow.log_metric("accuracy", accuracy)
        mlflow.sklearn.log_model(model, "model")
        
        results.append({param_dict, "accuracy": accuracy})

# 分析结果
df = pd.DataFrame(results)
best = df.loc[df["accuracy"].idxmax()]
print(f"最佳配置: {best.to_dict()}")

# 完整的实验记录模板
import mlflow
import json
from datetime import datetime

def log_complete_experiment(model, params, metrics, X_train, tags=None):
    with mlflow.start_run() as run:
        # 基础参数
        for k, v in params.items():
            mlflow.log_param(k, v)
        
        # 所有指标
        for k, v in metrics.items():
            mlflow.log_metric(k, v)
        
        # 保存模型
        mlflow.sklearn.log_model(model, "model")
        
        # 记录环境信息
        env_info = {
            "python_version": "3.10",
            "timestamp": datetime.now().isoformat(),
            "dataset_hash": "abc123",  # 来自 DVC
            "git_commit": "def456",   # 来自 Git
        }
        mlflow.log_dict(env_info, "env_info.json")
        
        # 设置标签
        if tags:
            for k, v in tags.items():
                mlflow.set_tag(k, v)
        
        return run.info.run_id

# 使用示例
run_id = log_complete_experiment(
    model=trained_model,
    params={"n_estimators": 100, "max_depth": 10},
    metrics={"accuracy": 0.95, "f1": 0.93},
    X_train=X_train,
    tags={"status": "success", "author": "alice"}
)

# 实验对比分析
import mlflow
import pandas as pd

# 搜索所有成功的实验
runs = mlflow.search_runs(
    experiment_names=["customer-churn"],
    filter_string="tags.status = 'success'",
    order_by=["metrics.f1 DESC"],
)

# 分析超参数影响
import seaborn as sns
import matplotlib.pyplot as plt

sns.scatterplot(
    data=runs,
    x="params.n_estimators",
    y="metrics.accuracy",
    hue="params.max_depth",
    size="params.min_samples_split",
)
plt.title("超参数对准确率的影响")
plt.savefig("hyperparameter_analysis.png")
mlflow.log_artifact("hyperparameter_analysis.png")

# 完整实验管理流程
import mlflow
import wandb
import dvc.api
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import yaml

# 1. 加载版本化数据
with dvc.api.open("data/train.csv", repo="git@github.com:org/project.git") as f:
    import pandas as pd
    df = pd.read_csv(f)

# 2. 加载实验配置
with open("config/experiment.yaml") as f:
    config = yaml.safe_load(f)

# 3. 同时使用 MLflow 和 WandB
mlflow.set_experiment(config["experiment_name"])
wandb.init(project=config["experiment_name"], config=config["params"])

X_train, X_test, y_train, y_test = train_test_split(
    df.drop("churn", axis=1), df["churn"], test_size=0.2
)

with mlflow.start_run() as run:
    # 记录参数到两个系统
    for k, v in config["params"].items():
        mlflow.log_param(k, v)
        wandb.config[k] = v
    
    # 训练
    model = RandomForestClassifier(**config["params"])
    model.fit(X_train, y_train)
    
    # 评估
    report = classification_report(y_test, model.predict(X_test), output_dict=True)
    for metric, value in report["accuracy"].items():
        if isinstance(value, (int, float)):
            mlflow.log_metric(metric, value)
            wandb.log({metric: value})
    
    # 保存
    mlflow.sklearn.log_model(model, "model")
    wandb.sklearn.plot_classifier(model, X_train, X_test, y_train, y_test)
    
    print(f"Run ID: {run.info.run_id}")

# config/experiment.yaml
experiment_name: "customer-churn-prediction"

params:
  n_estimators: 200
  max_depth: 15
  min_samples_split: 5
  random_state: 42

data:
  train: "data/train.csv"
  test: "data/test.csv"
  version: "v2.0"

tracking:
  mlflow:
    uri: "sqlite:///mlflow.db"
    artifact_root: "s3://my-bucket/artifacts"
  wandb:
    entity: "my-team"
    project: "customer-churn"