📚 课程目录第4章实验跟踪工具

第四章：实验跟踪工具——MLflow、Weights & Biases、Neptune

说实话，我刚开始做 AI 训练那会儿，实验管理基本靠「本地记事本 + 文件夹命名」。跑了一百多个实验后，自己都分不清哪个 config 对应哪个 checkpoint。后来一个前辈跟我说：别省那几分钟，用个实验跟踪工具，省下的时间够你睡好几觉。

嗯，这话我记住了。今天咱们就聊聊三个主流工具：MLflow、Weights & Biases（W&B）和 Neptune。它们都能帮你记录参数、指标、模型版本，但各有脾气。

为什么需要实验跟踪？

你想想看，一个典型训练实验有多少变量？学习率、批次大小、优化器、数据增强方式、网络层数……我见过一个团队，光学习率就试了 30 个不同值。如果不记录，结果出来根本不知道哪个参数起效。

实验跟踪的核心任务就三条：

记录——把每次实验的参数、代码版本、环境依赖、训练指标都存下来
对比——不同实验之间能快速看曲线、表格，找出最优组合
复现——拿到一个实验 ID，就能原样复现当初的结果

下面这张图能帮你快速理解这三个工具的定位差异。

MLflow：开源派的首选

MLflow 由 Databricks 打造，2018 年就开源了。它的理念很简单：轻量、开放、语言无关。你可以在本地跑，也可以在云上搭，不收一分钱。

我最早在项目中用 MLflow 是因为客户要求所有实验必须在内网完成，不能把数据传到第三方。MLflow 的 Tracking Server 自己搭，代码里加几行就搞定。

核心组件

Tracking —— 记录参数、指标、模型、工件
Projects —— 将代码打包成可复用的运行环境
Models —— 模型注册与推理服务
Registry —— 模型版本管理与阶段标注

看个最简例子：

import mlflow

# 开启一次实验
with mlflow.start_run():
    # 记录参数
    mlflow.log_param("learning_rate", 0.001)
    mlflow.log_param("batch_size", 64)
    
    # 模拟训练循环
    for epoch in range(10):
        loss = 0.5 ** (epoch + 1)
        mlflow.log_metric("loss", loss, step=epoch)
    
    # 保存模型
    mlflow.pytorch.log_model(model, "model")
    
    # 记录额外文件
    mlflow.log_artifact("config.yaml")

我个人习惯把模型注册到 MLflow Registry，然后用 mlflow.models.serve 一键启动 REST API。这在小团队里非常方便——不需要额外搭推理框架。

小技巧：如果你的实验在 Docker 容器里跑，记得挂载 mlruns 目录到宿主机，否则容器销毁后实验记录就丢了。

Weights & Biases：协作与可视化之王

如果你不介意把数据放到云端（有自托管方案），W&B 在可视化方面可以说是碾压级别的。它的实验面板可以同时对比几十个实验，曲线实时刷新，还能自动帮你做超参重要性分析。

第一次用 W&B 时，我跑了 20 个超参组合，它自动生成了平行坐标图和相关性热图，我一眼就看出「学习率 1e-4 配合 AdamW 效果最好」。这在我过去手动对比时要花大半天。

接入方式几乎一样简单：

import wandb

# 初始化，自动接手超参和代码版本
wandb.init(project="my-awesome-project", config={
    "learning_rate": 0.001,
    "batch_size": 64,
    "epochs": 10,
})

# 训练循环
for epoch in range(10):
    loss = 0.5 ** (epoch + 1)
    # 记录指标，自动生成图表
    wandb.log({"loss": loss, "accuracy": 1 - loss})
    
    # 偶尔记录图像/直方图
    if epoch % 3 == 0:
        wandb.log({"predictions": wandb.Image(pred_img)})

# 结束运行
wandb.finish()

注意：W&B 免费版有项目数和存储限制。我们团队一次百万级图片训练，日志膨胀到几十 GB，不得不升级到专业版。如果预算敏感，提前评估日志量。

W&B 还有一个杀手级功能：Sweeps。你只需要定义一个搜索空间配置：

program: train.py
method: bayes
metric:
  name: val_loss
  goal: minimize
parameters:
  learning_rate:
    min: 0.0001
    max: 0.01
    distribution: log_uniform
  dropout:
    values: [0.1, 0.2, 0.3, 0.5]
  optimizer:
    values: ["adam", "sgd"]

然后一个命令 wandb sweep <config.yaml>，它会自动调度多个 agent 并行搜索最优超参。我用它一次跑了 80 个实验，中间去睡了个觉，第二天就有了最优组合。

Neptune：结构化实验管理的瑞士军刀

Neptune 比较晚进入市场，但它在「结构化」方面做得特别细致。它把实验记录看作一个元数据图谱，你可以给实验打标签、添加自定义属性、关联 Git commit、甚至跟踪数据集的版本。

我前年做一个图像分割项目，需要同时跟踪输入图片的分辨率、GPU 显存使用率、模型参数量、和每个类别的 IoU。Neptune 的自定义仪表盘让我把这些信息都放在一个页面，客户看了直呼专业。

基本用法：

import neptune

# 创建运行实例
run = neptune.init_run(
    project="my-team/my-project",
    api_token="your_api_token",
)

# 记录参数
run["hyperparameters"] = {
    "lr": 0.001,
    "batch_size": 32,
    "model": "UNet"
}

# 训练循环
for epoch in range(10):
    loss = 0.5 ** (epoch + 1)
    run["metrics/train/loss"].append(loss)
    run["metrics/val/iou"].append(0.8 - 0.1 * epoch)

# 记录代码版本
run["source_code"].upload("train.py")

run.stop()

Neptune 的强项在于可查询。你可以用类似 SQL 的语法：run["hyperparameters/lr"] > 0.001 and run["metrics/val/iou"] > 0.7 来筛选实验。我曾经用这个功能从 500 个实验中筛出了所有「大学习率 + 低 dropout」的组合，帮团队快速定位问题。

横向对比：怎么选？

维度	MLflow	Weights & Biases	Neptune
部署方式	完全自托管 or 托管版	云端为主，有自托管版	云端为主，有私有云版
可视化强度	★★☆☆☆ (基本曲线)	★★★★★ (交互面板)	★★★★☆ (自定义仪表盘)
超参调优	需配合 Optuna	内置 Sweeps	内置阈值/搜索
模型注册	内置 Registry	Model Registry (付费)	内置 Registry
团队协作	通过共享 SQL 数据库	原生协作 + 报告共享	原生团队 + 权限管理
代码/数据溯源	需手动记录 commit	自动记录 Git + DVC 集成	自动 Git + 数据版本快照
开源	是 (Apache 2.0)	SDK 开源，后端闭源	SDK 开源，后端闭源
启动速度	极快 (纯 Python 无网络)	较快 (需网络验证)	较慢 (初始化检查多)

我的选择策略

说实话，没有万能工具。我一般按场景来：

本地研究、快速验证 → MLflow。它够轻，不会干扰训练流程，而且完全离线。
团队协作、需要汇报 → W&B。它的报告功能能直接把实验曲线嵌入到网页里，老板看了直点头。
企业级合规、多维度元数据管理 → Neptune。它能把模型、数据、代码、指标都关联成图，审计时一查一个准。

核心原则：无论选哪个，第一件事是统一接口。我在团队里强制要求所有训练脚本必须同时记录参数、指标和模型文件，哪怕只跑一次实验。这样后续做版本锁定和复现时，才能从工具里直接拉出完整上下文。

版本锁定与实验记录的联动

别忘了，我们这章的主题是「实验跟踪工具」，它的终极目标是服务于版本锁定与复现。所以，我建议在每个实验记录里必须包含以下信息（无论用哪个工具）：

# 理想实验记录结构
experiment:
  id: "exp_20250321_1245"
  timestamp: 2025-03-21 12:45:00
  git_commit: "a3b2c1d"
  dataset_version: "v2.1"
  environment:
    python: "3.10.12"
    torch: "2.1.0"
    cuda: "12.1"
  parameters:
    learning_rate: 0.001
    batch_size: 64
    optimizer: "AdamW"
  metrics:
    train_loss: 0.023
    val_accuracy: 0.971
  artifacts:
    model: "s3://bucket/v1/model.pt"
    config: "configs/exp_1245.yaml"
    logs: "runs/exp_1245/tensorboard"

你想想看，半年后你要复现这个实验，只需要拿出实验 ID，工具会自动拉出当时的所有参数、代码版本、依赖环境。这才是真正意义上的「一键复现」。

经验分享：我曾经接到一个紧急任务——复现三个月前客户认可的 baseline。客户报了个实验编号，我用 W&B 直接恢复了当时的 config 和模型，半小时内跑出了完全一致的精度曲线。老板愣了半天，说「这工具钱花得值」。

好了，实验跟踪工具就讲到这里。记住一句话：记录越完整，复现越轻松。别偷懒，每次实验都多花三十秒把参数和指标写进工具，将来你会感谢自己的。