实验管理

返回工程实践

训练模型时往往需要尝试数十种超参数组合,实验管理工具自动记录每次实验的配置、指标、产物,让结果可复现、可对比。


核心概念

概念说明
Run(运行)一次独立的训练/评估过程
Experiment(实验)相关 Run 的集合(如同一任务的不同配置)
Artifact(产物)模型权重、数据集、评估报告等文件
Metric(指标)随步骤变化的数值(loss、accuracy 等)
Parameter(参数)固定的超参数配置

MLflow

开源、自托管,框架无关:

import mlflow
import mlflow.pytorch
 
mlflow.set_experiment("bert-text-classification")
 
with mlflow.start_run(run_name="lr-3e-5-batch-32"):
    # 记录超参数
    mlflow.log_params({
        "learning_rate": 3e-5,
        "batch_size": 32,
        "epochs": 10,
        "model": "bert-base-chinese",
    })
 
    for epoch in range(10):
        train_loss, val_acc = train_epoch(model, ...)
 
        # 记录指标(随 step 变化)
        mlflow.log_metrics({
            "train_loss": train_loss,
            "val_accuracy": val_acc,
        }, step=epoch)
 
    # 保存模型
    mlflow.pytorch.log_model(model, "model")
 
    # 记录任意文件
    mlflow.log_artifact("confusion_matrix.png")
    mlflow.log_artifact("config.yaml")
# 启动 UI
mlflow ui --host 0.0.0.0 --port 5000

MLflow 模型注册中心:

# 将模型提升到生产
from mlflow.tracking import MlflowClient
 
client = MlflowClient()
# 注册模型
model_uri = f"runs:/{run_id}/model"
mv = mlflow.register_model(model_uri, "TextClassifier")
 
# 设置阶段:Staging → Production
client.transition_model_version_stage(
    name="TextClassifier",
    version=mv.version,
    stage="Production",
)

Weights & Biases(W&B)

云端,可视化丰富,团队协作强:

import wandb
 
wandb.init(
    project="llm-finetuning",
    name="qwen-lora-r16",
    config={
        "model": "Qwen2.5-7B",
        "lora_r": 16,
        "learning_rate": 2e-4,
        "batch_size": 4,
    }
)
 
for step, batch in enumerate(dataloader):
    loss = train_step(batch)
 
    wandb.log({
        "train/loss": loss,
        "train/grad_norm": grad_norm,
        "lr": scheduler.get_last_lr()[0],
    }, step=step)
 
# 保存模型到 W&B Artifacts
artifact = wandb.Artifact("qwen-lora-r16", type="model")
artifact.add_dir("./output")
wandb.log_artifact(artifact)
 
wandb.finish()

HuggingFace Trainer 集成(自动记录):

from transformers import TrainingArguments
 
args = TrainingArguments(
    output_dir="./output",
    report_to="wandb",  # 自动记录所有指标
    run_name="qwen-lora-experiment",
    ...
)

超参数搜索

Optuna(贝叶斯优化)

import optuna
 
def objective(trial):
    lr = trial.suggest_float("lr", 1e-5, 1e-3, log=True)
    batch_size = trial.suggest_categorical("batch_size", [16, 32, 64])
    dropout = trial.suggest_float("dropout", 0.1, 0.5)
 
    model = build_model(dropout=dropout)
    val_acc = train_and_evaluate(model, lr=lr, batch_size=batch_size)
    return val_acc
 
study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=50)
 
print("最优参数:", study.best_params)
print("最优值:", study.best_value)

W&B Sweep

# sweep.yaml
program: train.py
method: bayes   # random / grid / bayes
metric:
  goal: maximize
  name: val/accuracy
parameters:
  learning_rate:
    distribution: log_uniform_values
    min: 1e-5
    max: 1e-3
  batch_size:
    values: [16, 32, 64]
  dropout:
    distribution: uniform
    min: 0.1
    max: 0.5
wandb sweep sweep.yaml      # 创建 sweep,返回 SWEEP_ID
wandb agent SWEEP_ID        # 启动 Agent 自动跑实验

实验对比最佳实践

1. 每次只改一个变量(控制变量法)
2. 固定随机种子(torch.manual_seed / np.random.seed)
3. 同一超参数跑 3 次取均值(消除随机性影响)
4. 记录环境信息(库版本、GPU 型号、数据集版本)
5. 给 Run 起有意义的名字(lr-3e5_batch-32_epoch-10)
6. 失败的实验也要记录原因(避免重复踩坑)

工具对比

维度MLflowW&BTensorBoard
部署自托管云端本地
模型注册
超参数搜索❌(集成 Optuna)✅ Sweep
团队协作一般
免费额度完全免费个人免费完全免费
LLM 支持一般强(Prompts 追踪)

相关文档