实验管理
→ 返回工程实践
训练模型时往往需要尝试数十种超参数组合,实验管理工具自动记录每次实验的配置、指标、产物,让结果可复现、可对比。
核心概念
| 概念 | 说明 |
|---|---|
| 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 5000MLflow 模型注册中心:
# 将模型提升到生产
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.5wandb 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. 失败的实验也要记录原因(避免重复踩坑)
工具对比
| 维度 | MLflow | W&B | TensorBoard |
|---|---|---|---|
| 部署 | 自托管 | 云端 | 本地 |
| 模型注册 | ✅ | ✅ | ❌ |
| 超参数搜索 | ❌(集成 Optuna) | ✅ Sweep | ❌ |
| 团队协作 | 一般 | 强 | 弱 |
| 免费额度 | 完全免费 | 个人免费 | 完全免费 |
| LLM 支持 | 一般 | 强(Prompts 追踪) | 弱 |