欢迎回来

登录 EAKE AI,继续您的智能之旅

忘记密码?
还没有账号?立即注册

W&B 实验追踪

W&B 实验追踪

W&B:记录 ML 实验、超参扫描、模型注册、仪表盘

Weights & Biases: ML 实验追踪 & MLOps

使用场景 This Skill

当需要以下场景时使用 Weights & Biases(W&B):

  • 追踪 ML 实验,自动记录指标
  • 实时可视化训练 仪表盘
  • 跨超参和配置比较运行
  • 自动超参优化 sweeps
  • 模型注册管理,含版本和 lineage
  • ML 项目团队协作
  • 追踪数据集、模型、代码等产物 并保留 lineage
  • 用户: 200,000+ ML 从业者 | GitHub Stars: 10.5k+ | 集成: 100+

    安装

    
    # 安装 W&B
    pip install wandb
    
    # 登录(创建 API key)
    wandb login
    
    # 或通过代码设置 API key
    export WANDB_API_KEY=your_api_key_here
    

    快速入门

    基础实验追踪

    
    import wandb
    
    # 初始化一个 run
    run = wandb.init(
        project="my-project",
        config={
            "learning_rate": 0.001,
            "epochs": 10,
            "batch_size": 32,
            "architecture": "ResNet50"
        }
    )
    
    # 训练循环
    for epoch in range(run.config.epochs):
        # 你的训练代码
        train_loss = train_epoch()
        val_loss = validate()
    
        # 记录指标
        wandb.log({
            "epoch": epoch,
            "train/loss": train_loss,
            "val/loss": val_loss,
            "train/accuracy": train_acc,
            "val/accuracy": val_acc
        })
    
    # 结束 run
    run.finish()
    

    PyTorch 集成

    
    import torch
    import wandb
    
    # 初始化
    wandb.init(project="pytorch-demo", config={
        "lr": 0.001,
        "epochs": 10
    })
    
    # 访问 config
    config = wandb.config
    
    # 训练循环
    for epoch in range(config.epochs):
        for batch_idx, (data, target) in enumerate(train_loader):
            # 前向传播
            output = model(data)
            loss = criterion(output, target)
    
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
            # 每 100 个 batch 记录一次
            if batch_idx % 100 == 0:
                wandb.log({
                    "loss": loss.item(),
                    "epoch": epoch,
                    "batch": batch_idx
                })
    
    # 保存模型
    torch.save(model.state_dict(), "model.pth")
    wandb.save("model.pth")  # 上传到 W&B
    
    wandb.finish()
    

    核心概念

    1. Projects 和 Runs

    Project:相关实验的集合

    Run:训练脚本的一次执行

    
    # 创建/使用 project
    run = wandb.init(
        project="image-classification",
        name="resnet50-experiment-1",  # 可选 run 名称
        tags=["baseline", "resnet"],    # 用标签组织
        notes="First baseline run"      # 添加备注
    )
    
    # 每个 run 有唯一 ID
    print(f"Run ID: {run.id}")
    print(f"Run URL: {run.url}")
    

    2. 配置追踪

    自动追踪超参数:

    
    config = {
        # 模型架构
        "model": "ResNet50",
        "pretrained": True,
    
        # 训练参数
        "learning_rate": 0.001,
        "batch_size": 32,
        "epochs": 50,
        "optimizer": "Adam",
    
        # 数据参数
        "dataset": "ImageNet",
        "augmentation": "standard"
    }
    
    wandb.init(project="my-project", config=config)
    
    # 训练期间访问 config
    lr = wandb.config.learning_rate
    batch_size = wandb.config.batch_size
    

    3. 指标记录

    
    # 记录标量
    wandb.log({"loss": 0.5, "accuracy": 0.92})
    
    # 记录多个指标
    wandb.log({
        "train/loss": train_loss,
        "train/accuracy": train_acc,
        "val/loss": val_loss,
        "val/accuracy": val_acc,
        "learning_rate": current_lr,
        "epoch": epoch
    })
    
    # 用自定义 x 轴记录
    wandb.log({"loss": loss}, step=global_step)
    
    # 记录媒体(图像、音频、视频)
    wandb.log({"examples": [wandb.Image(img) for img in images]})
    
    # 记录直方图
    wandb.log({"gradients": wandb.Histogram(gradients)})
    
    # 记录表格
    table = wandb.Table(columns=["id", "prediction", "ground_truth"])
    wandb.log({"predictions": table})
    

    4. 模型检查点

    
    import torch
    import wandb
    
    # 保存模型检查点
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
    }
    
    torch.save(checkpoint, 'checkpoint.pth')
    
    # 上传到 W&B
    wandb.save('checkpoint.pth')
    
    # 或使用 Artifacts(推荐)
    artifact = wandb.Artifact('model', type='model')
    artifact.add_file('checkpoint.pth')
    wandb.log_artifact(artifact)
    

    超参数 Sweeps

    自动搜索最优超参数。

    定义 Sweep 配置

    
    sweep_config = {
        'method': 'bayes',  # 或 'grid'、'random'
        'metric': {
            'name': 'val/accuracy',
            'goal': 'maximize'
        },
        'parameters': {
            'learning_rate': {
                'distribution': 'log_uniform',
                'min': 1e-5,
                'max': 1e-1
            },
            'batch_size': {
                'values': [16, 32, 64, 128]
            },
            'optimizer': {
                'values': ['adam', 'sgd', 'rmsprop']
            },
            'dropout': {
                'distribution': 'uniform',
                'min': 0.1,
                'max': 0.5
            }
        }
    }
    
    # 初始化 sweep
    sweep_id = wandb.sweep(sweep_config, project="my-project")
    

    定义训练函数

    
    def train():
        # 初始化 run
        run = wandb.init()
    
        # 访问 sweep 参数
        lr = wandb.config.learning_rate
        batch_size = wandb.config.batch_size
        optimizer_name = wandb.config.optimizer
    
        # 根据 sweep config 构建模型
        model = build_model(wandb.config)
        optimizer = get_optimizer(optimizer_name, lr)
    
        # 训练循环
        for epoch in range(NUM_EPOCHS):
            train_loss = train_epoch(model, optimizer, batch_size)
            val_acc = validate(model)
    
            # 记录指标
            wandb.log({
                "train/loss": train_loss,
                "val/accuracy": val_acc
            })
    
    # 运行 sweep
    wandb.agent(sweep_id, function=train, count=50)  # 运行 50 次试验
    

    Sweep 策略

    
    # 网格搜索 - 穷举
    sweep_config = {
        'method': 'grid',
        'parameters': {
            'lr': {'values': [0.001, 0.01, 0.1]},
            'batch_size': {'values': [16, 32, 64]}
        }
    }
    
    # 随机搜索
    sweep_config = {
        'method': 'random',
        'parameters': {
            'lr': {'distribution': 'uniform', 'min': 0.0001, 'max': 0.1},
            'dropout': {'distribution': 'uniform', 'min': 0.1, 'max': 0.5}
        }
    }
    
    # 贝叶斯优化(推荐)
    sweep_config = {
        'method': 'bayes',
        'metric': {'name': 'val/loss', 'goal': 'minimize'},
        'parameters': {
            'lr': {'distribution': 'log_uniform', 'min': 1e-5, 'max': 1e-1}
        }
    }
    

    Artifacts

    追踪数据集、模型及其他文件的 lineage。

    记录 Artifacts

    
    # 创建 artifact
    artifact = wandb.Artifact(
        name='training-dataset',
        type='dataset',
        description='ImageNet training split',
        metadata={'size': '1.2M images', 'split': 'train'}
    )
    
    # 添加文件
    artifact.add_file('data/train.csv')
    artifact.add_dir('data/images/')
    
    # 记录 artifact
    wandb.log_artifact(artifact)
    

    使用 Artifacts

    
    # 下载并使用 artifact
    run = wandb.init(project="my-project")
    
    # 下载 artifact
    artifact = run.use_artifact('training-dataset:latest')
    artifact_dir = artifact.download()
    
    # 使用数据
    data = load_data(f"{artifact_dir}/train.csv")
    

    模型注册表

    
    # 将模型记录为 artifact
    model_artifact = wandb.Artifact(
        name='resnet50-model',
        type='model',
        metadata={'architecture': 'ResNet50', 'accuracy': 0.95}
    )
    
    model_artifact.add_file('model.pth')
    wandb.log_artifact(model_artifact, aliases=['best', 'production'])
    
    # 链接到模型注册表
    run.link_artifact(model_artifact, 'model-registry/production-models')
    

    集成示例

    HuggingFace Transformers

    
    from transformers import Trainer, TrainingArguments
    import wandb
    
    # 初始化 W&B
    wandb.init(project="hf-transformers")
    
    # 训练参数含 W&B
    training_args = TrainingArguments(
        output_dir="./results",
        report_to="wandb",  # 启用 W&B 日志
        run_name="bert-finetuning",
        logging_steps=100,
        save_steps=500
    )
    
    # Trainer 自动记录到 W&B
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset
    )
    
    trainer.train()
    

    PyTorch Lightning

    
    from pytorch_lightning import Trainer
    from pytorch_lightning.loggers import WandbLogger
    import wandb
    
    # 创建 W&B logger
    wandb_logger = WandbLogger(
        project="lightning-demo",
        log_model=True  # 记录模型检查点
    )
    
    # 与 Trainer 配合使用
    trainer = Trainer(
        logger=wandb_logger,
        max_epochs=10
    )
    
    trainer.fit(model, datamodule=dm)
    

    Keras/TensorFlow

    
    import wandb
    from wandb.keras import WandbCallback
    
    # 初始化
    wandb.init(project="keras-demo")
    
    # 添加 callback
    model.fit(
        x_train, y_train,
        validation_data=(x_val, y_val),
        epochs=10,
        callbacks=[WandbCallback()]  # 自动记录指标
    )
    

    可视化与分析

    自定义图表

    
    # 记录自定义可视化
    import matplotlib.pyplot as plt
    
    fig, ax = plt.subplots()
    ax.plot(x, y)
    wandb.log({"custom_plot": wandb.Image(fig)})
    
    # 记录混淆矩阵
    wandb.log({"conf_mat": wandb.plot.confusion_matrix(
        probs=None,
        y_true=ground_truth,
        preds=predictions,
        class_names=class_names
    )})
    

    报告

    在 W&B UI 中创建可共享报告:

  • 组合 runs、图表和文本
  • 支持 Markdown
  • 可嵌入可视化
  • 团队协作
  • 最佳实践

    1. 用 Tags 和 Groups 组织

    
    wandb.init(
        project="my-project",
        tags=["baseline", "resnet50", "imagenet"],
        group="resnet-experiments",  # 将相关 runs 归组
        job_type="train"             # 任务类型
    )
    

    2. 记录所有相关内容

    
    # 记录系统指标
    wandb.log({
        "gpu/util": gpu_utilization,
        "gpu/memory": gpu_memory_used,
        "cpu/util": cpu_utilization
    })
    
    # 记录代码版本
    wandb.log({"git_commit": git_commit_hash})
    
    # 记录数据划分
    wandb.log({
        "data/train_size": len(train_dataset),
        "data/val_size": len(val_dataset)
    })
    

    3. 使用描述性名称

    
    # ✅ 好:描述性 run 名称
    wandb.init(
        project="nlp-classification",
        name="bert-base-lr0.001-bs32-epoch10"
    )
    
    # ❌ 差:通用名称
    wandb.init(project="nlp", name="run1")
    

    4. 保存重要 Artifacts

    
    # 保存最终模型
    artifact = wandb.Artifact('final-model', type='model')
    artifact.add_file('model.pth')
    wandb.log_artifact(artifact)
    
    # 保存预测结果用于分析
    predictions_table = wandb.Table(
        columns=["id", "input", "prediction", "ground_truth"],
        data=predictions_data
    )
    wandb.log({"predictions": predictions_table})
    

    5. 不稳定网络用离线模式

    
    import os
    
    # 启用离线模式
    os.environ["WANDB_MODE"] = "offline"
    
    wandb.init(project="my-project")
    # ... 你的代码 ...
    
    # 稍后同步
    # wandb sync 
    

    团队协作

    共享 Runs

    
    # Runs 自动可通过 URL 共享
    run = wandb.init(project="team-project")
    print(f"分享此 URL: {run.url}")
    

    团队项目

  • 在 wandb.ai 创建团队账户
  • 添加团队成员
  • 设置项目可见性(私有/公开)
  • 使用团队级 artifacts 和模型注册表
  • 定价

  • 免费:无限公开项目,100GB 存储
  • 学术:学生/研究人员免费
  • 团队:$50/人/月,私有项目,无限存储
  • 企业:定制定价,支持本地部署
  • 资源

  • 文档: https://docs.wandb.ai
  • GitHub: https://github.com/wandb/wandb (10.5k+ stars)
  • 示例: https://github.com/wandb/examples
  • 社区: https://wandb.ai/community
  • Discord: https://wandb.me/discord
  • 另见

  • references/sweeps.md - 全面超参数优化指南
  • references/artifacts.md - 数据和模型版本管理模式
  • references/integrations.md - 框架特定示例
  • 安装指南

    复制下方命令,在终端运行即可安装:

    # 安装到当前项目
    npx skills add weights-and-biases
    # 全局安装 — 所有项目可用
    npx skills add weights-and-biases -g

    使用指南

    安装完成后,在对话框中直接使用此技能。

    基本信息
    作者 Community 分类 coding 难度 Intermediate 时长 1 hour
    🛠️ 安装命令
    # 安装到当前项目
    npx skills add weights-and-biases
    # 全局安装
    npx skills add weights-and-biases -g

    发表评论