Skip to content

Weights & Biases (W&B)

Weights & Biases（简称 W&B 或 wandb）是一个实验追踪与可视化平台，专为机器学习团队设计。相比 TensorBoard，W&B 提供了云端存储、团队协作、自动超参数搜索等高级功能。

基本使用

安装与配置

pip install wandb
wandb login  # 输入 API Key（从 wandb.ai 获取）

基础集成

import wandb

# 初始化实验
wandb.init(
    project="my-project",
    name="resnet50-baseline",
    config={
        "learning_rate": 1e-3,
        "batch_size": 32,
        "epochs": 100,
        "architecture": "ResNet-50",
        "optimizer": "AdamW",
    }
)

for epoch in range(num_epochs):
    train_loss = train_one_epoch(model, train_loader, optimizer)
    val_loss, val_acc = evaluate(model, val_loader)

    # 记录指标
    wandb.log({
        "train_loss": train_loss,
        "val_loss": val_loss,
        "val_accuracy": val_acc,
        "lr": optimizer.param_groups[0]['lr'],
        "epoch": epoch,
    })

# 保存模型
wandb.save("model_best.pth")
wandb.finish()

核心功能

实验对比

W&B 自动为每次 wandb.init() 创建一个 Run。在 Web 界面上可以：

  • 并排对比不同实验的训练曲线
  • 按超参数筛选和排序实验
  • 创建自定义 Dashboard

超参数搜索 (Sweep)

W&B Sweep 提供自动化的超参数搜索：

# 定义搜索空间
sweep_config = {
    "method": "bayes",  # bayes, grid, random
    "metric": {"name": "val_accuracy", "goal": "maximize"},
    "parameters": {
        "learning_rate": {"min": 1e-5, "max": 1e-2, "distribution": "log_uniform_values"},
        "batch_size": {"values": [16, 32, 64, 128]},
        "weight_decay": {"min": 1e-5, "max": 1e-1, "distribution": "log_uniform_values"},
    }
}

sweep_id = wandb.sweep(sweep_config, project="my-project")

def train():
    wandb.init()
    config = wandb.config
    # 使用 config.learning_rate, config.batch_size 等进行训练
    ...

wandb.agent(sweep_id, function=train, count=50)  # 运行50次实验

Artifacts（制品管理）

用于版本化管理数据集、模型和其他文件：

# 保存模型为 artifact
artifact = wandb.Artifact('trained-model', type='model')
artifact.add_file('model_best.pth')
wandb.log_artifact(artifact)

# 加载 artifact
artifact = wandb.use_artifact('trained-model:latest')
artifact_dir = artifact.download()

表格与可视化

# 记录预测样本
table = wandb.Table(columns=["image", "prediction", "ground_truth"])
for img, pred, gt in samples:
    table.add_data(wandb.Image(img), pred, gt)
wandb.log({"predictions": table})

与其他工具对比

功能 TensorBoard W&B MLflow
实验追踪 本地 云端 本地/远程
团队协作 不支持 原生支持 支持
超参数搜索 不支持 Sweep（贝叶斯等） 不支持（需集成）
模型注册 不支持 Artifacts Model Registry
免费方案 完全免费 个人免费 开源免费
学习成本 低 低 中等

选型建议：

  • 个人项目/快速实验 → TensorBoard（零配置）
  • 团队协作/正式项目 → W&B（功能最全面）
  • 需要自托管/MLOps Pipeline → MLflow

参考

评论 #