Weights & Biases(W&B)实验管理实战指南：从实验追踪到 Model Registry 与生产环境监控

┌─────────────────────────────────────────────────────────────────┐
│                    W&B Platform Architecture                     │
├─────────────────────────────────────────────────────────────────┤
│                                                                   │
│  ┌──────────────┐   ┌──────────────┐   ┌──────────────────────┐  │
│  │  Experiment   │   │   Sweeps     │   │     Artifacts        │  │
│  │  Tracking     │──▶│ Hyperparams  │──▶│  Data/Model Version  │  │
│  │  (wandb.log)  │   │  Tuning      │   │  Control             │  │
│  └──────┬───────┘   └──────────────┘   └──────────┬───────────┘  │
│         │                                          │              │
│         ▼                                          ▼              │
│  ┌──────────────┐   ┌──────────────┐   ┌──────────────────────┐  │
│  │   Tables      │   │   Reports    │   │   Model Registry     │  │
│  │  Data Viz     │   │  Team Collab │   │  Production Staging  │  │
│  │  Analysis     │   │  Sharing     │   │  Deployment          │  │
│  └──────────────┘   └──────────────┘   └──────────┬───────────┘  │
│                                                    │              │
│                                          ┌─────────▼──────────┐  │
│                                          │   W&B Launch /      │  │
│                                          │   Webhooks / CI/CD  │  │
│                                          │   Automation        │  │
│                                          └────────────────────┘  │
│                                                                   │
│  Infrastructure: Cloud (SaaS) / Self-Hosted / Dedicated Cloud    │
└─────────────────────────────────────────────────────────────────┘

import wandb
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

# 1. 初始化 W&B —— 指定项目与实验名称
wandb.init(
    project="image-classifier",
    name="resnet50-lr0.001-batch32",
    config={
        "architecture": "ResNet50",
        "learning_rate": 0.001,
        "batch_size": 32,
        "epochs": 100,
        "optimizer": "AdamW",
        "weight_decay": 1e-4,
        "scheduler": "CosineAnnealingLR",
        "dataset": "CIFAR-100",
        "augmentation": "RandAugment",
    },
    tags=["baseline", "resnet", "cifar100"],
    notes="ResNet50 baseline with CosineAnnealing scheduler",
)

config = wandb.config  # 通过 config 对象访问超参数

# 2. 设置模型、优化器、调度器
model = build_model(config.architecture)
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=config.learning_rate,
    weight_decay=config.weight_decay
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=config.epochs
)

# 3. 在 W&B 中追踪模型梯度(可选)
wandb.watch(model, log="all", log_freq=100)

# 4. 训练循环 —— 记录指标
for epoch in range(config.epochs):
    model.train()
    train_loss = 0.0
    correct = 0
    total = 0

    for batch_idx, (images, labels) in enumerate(train_loader):
        images, labels = images.cuda(), labels.cuda()
        outputs = model(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

        # 按 batch 记录(可选)
        if batch_idx % 50 == 0:
            wandb.log({
                "batch/loss": loss.item(),
                "batch/step": epoch * len(train_loader) + batch_idx,
            })

    # 按 epoch 记录指标
    train_acc = 100.0 * correct / total
    val_loss, val_acc = evaluate(model, val_loader)

    wandb.log({
        "epoch": epoch,
        "train/loss": train_loss / len(train_loader),
        "train/accuracy": train_acc,
        "val/loss": val_loss,
        "val/accuracy": val_acc,
        "learning_rate": scheduler.get_last_lr()[0],
        "gpu/memory_allocated_gb": torch.cuda.memory_allocated() / 1e9,
        "gpu/utilization": torch.cuda.utilization(),
    })

    scheduler.step()

    # 保存性能最佳的模型 checkpoint
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(model.state_dict(), "best_model.pth")
        wandb.save("best_model.pth")
        wandb.run.summary["best_val_accuracy"] = best_val_acc

# 5. 结束实验
wandb.finish()

# Hugging Face Transformers 集成示例
from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./results",
    report_to="wandb",           # 自动接入 W&B
    run_name="bert-finetune-v1",
    logging_steps=10,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="eval_f1",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)

trainer.train()  # 自动将指标记录到 W&B

import wandb

# 1. 定义 Sweep 配置
sweep_config = {
    "method": "bayes",  # 可选 grid、random、bayes
    "metric": {
        "name": "val/accuracy",
        "goal": "maximize"
    },
    "parameters": {
        "learning_rate": {
            "distribution": "log_uniform_values",
            "min": 1e-5,
            "max": 1e-2,
        },
        "batch_size": {
            "values": [16, 32, 64, 128]
        },
        "optimizer": {
            "values": ["adam", "adamw", "sgd"]
        },
        "weight_decay": {
            "distribution": "log_uniform_values",
            "min": 1e-6,
            "max": 1e-2,
        },
        "dropout": {
            "distribution": "uniform",
            "min": 0.1,
            "max": 0.5,
        },
        "architecture": {
            "values": ["resnet18", "resnet34", "resnet50"]
        },
    },
    # Early Termination：自动中止表现较差的实验
    "early_terminate": {
        "type": "hyperband",
        "min_iter": 5,       # 至少运行 5 个 epoch 后再评估
        "eta": 3,            # 只保留排名前 1/3 的实验继续运行
        "max_iter": 50,
    },
}

# 2. 创建 Sweep
sweep_id = wandb.sweep(
    sweep=sweep_config,
    project="image-classifier"
)

# 3. 定义训练函数
def train():
    wandb.init()
    config = wandb.config

    model = build_model(config.architecture, dropout=config.dropout)
    optimizer = get_optimizer(
        config.optimizer, model, config.learning_rate, config.weight_decay
    )

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

        wandb.log({
            "epoch": epoch,
            "train/loss": train_loss,
            "val/loss": val_loss,
            "val/accuracy": val_acc,
        })

    wandb.finish()

# 4. 运行 Sweep Agent —— 最多执行 50 次实验
wandb.agent(sweep_id, function=train, count=50)

# 终端 1(GPU 0)
CUDA_VISIBLE_DEVICES=0 wandb agent your-entity/image-classifier/sweep_id

# 终端 2(GPU 1)
CUDA_VISIBLE_DEVICES=1 wandb agent your-entity/image-classifier/sweep_id

# 终端 3(其他机器)
CUDA_VISIBLE_DEVICES=0 wandb agent your-entity/image-classifier/sweep_id

import wandb

# === 上传数据集 ===
run = wandb.init(project="image-classifier", job_type="data-pipeline")

# 创建 Artifact
dataset_artifact = wandb.Artifact(
    name="cifar100-processed",
    type="dataset",
    description="预处理完成的 CIFAR-100 数据集(已应用归一化、数据增强)",
    metadata={
        "num_classes": 100,
        "train_size": 50000,
        "val_size": 10000,
        "preprocessing": "normalize_mean_std",
        "augmentation": "randaugment_n2_m14",
        "source": "torchvision.datasets.CIFAR100",
    },
)

# 添加目录或文件
dataset_artifact.add_dir("./data/processed/cifar100")
dataset_artifact.add_file("./data/metadata.json")

# 上传 —— 自动进行 v0、v1、v2… 版本管理
run.log_artifact(dataset_artifact)
run.finish()

# === 训练时下载数据集 ===
run = wandb.init(project="image-classifier", job_type="training")

# 使用特定版本或最新版本
dataset = run.use_artifact("cifar100-processed:latest")
data_dir = dataset.download()  # 返回本地缓存路径

# 进行训练...
train_loader = create_dataloader(data_dir, split="train")

# === 保存模型 checkpoint ===
model_artifact = wandb.Artifact(
    name="resnet50-cifar100",
    type="model",
    description="ResNet50 trained on CIFAR-100",
    metadata={
        "val_accuracy": best_val_acc,
        "architecture": "resnet50",
        "framework": "pytorch",
        "training_run": wandb.run.id,
    },
)

model_artifact.add_file("best_model.pth")
model_artifact.add_file("model_config.json")
run.log_artifact(model_artifact)
run.finish()

import wandb

# 1. 训练完成后将模型注册到 Registry
run = wandb.init(project="image-classifier", job_type="register-model")

# 引用已有的模型 Artifact
model_artifact = run.use_artifact("resnet50-cifar100:v5")

# 链接到 Registry —— 创建 "Registered Model" 或为已有模型添加新版本
run.link_artifact(
    artifact=model_artifact,
    target_path="model-registry/image-classifier-prod",
    aliases=["candidate", "v2.1.0"],
)
run.finish()

# 2. 生产部署时从 Registry 获取模型
run = wandb.init(project="image-classifier", job_type="deployment")

# 从 Registry 按指定 alias 加载模型
model_artifact = run.use_artifact(
    "model-registry/image-classifier-prod:production"
)
model_dir = model_artifact.download()

# 加载并提供模型服务
model = load_model(model_dir)
run.finish()

# 3. 通过 Webhook 实现自动化 —— 响应 Registry 事件
# 可在 W&B UI 中设置，也可通过 API 配置
# 新模型获得 "production" alias 时自动触发部署流水线

import wandb

run = wandb.init(project="image-classifier", job_type="evaluation")

# 创建预测结果表
columns = ["image", "true_label", "predicted_label", "confidence", "correct"]
prediction_table = wandb.Table(columns=columns)

model.eval()
with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images.cuda())
        probs = torch.softmax(outputs, dim=1)
        confidences, predictions = probs.max(dim=1)

        for i in range(len(images)):
            is_correct = predictions[i].item() == labels[i].item()
            prediction_table.add_data(
                wandb.Image(images[i]),                    # 图像
                class_names[labels[i].item()],             # 真实标签
                class_names[predictions[i].item()],        # 预测标签
                confidences[i].item(),                     # 置信度
                is_correct,                                # 是否正确
            )

# 记录表格
run.log({"predictions": prediction_table})

# 误分类分析表
error_table = wandb.Table(
    columns=["image", "true", "predicted", "confidence", "loss"]
)
# ... 仅筛选误分类样本后添加

run.log({"error_analysis": error_table})
run.finish()

organization/
├── project-name-research/     # 探索性实验
├── project-name-staging/      # 验证候选部署模型
└── project-name-production/   # 生产环境监控

{模型架构}-{数据集}-{关键改动}-{日期}
例: resnet50-cifar100-cosine_lr-20260308

# .github/workflows/model-training.yml
name: Model Training and Evaluation Pipeline

on:
  push:
    paths:
      - 'src/models/**'
      - 'configs/**'
  workflow_dispatch:
    inputs:
      experiment_name:
        description: 'Experiment name'
        required: true

env:
  WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
  WANDB_PROJECT: image-classifier
  WANDB_ENTITY: your-team

jobs:
  train:
    runs-on: [self-hosted, gpu]
    steps:
      - uses: actions/checkout@v4

      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'

      - name: Install dependencies
        run: pip install -r requirements.txt

      - name: Train model
        run: |
          python src/train.py \
            --config configs/production.yaml \
            --wandb-run-name "ci-${{ github.sha }}"

      - name: Evaluate model
        run: |
          python src/evaluate.py \
            --model-artifact "resnet50-cifar100:latest" \
            --threshold-accuracy 0.92

      - name: Register model to W&B Registry
        if: success()
        run: |
          python src/register_model.py \
            --artifact "resnet50-cifar100:latest" \
            --registry "image-classifier-prod" \
            --alias "candidate"

  deploy:
    needs: train
    runs-on: ubuntu-latest
    if: github.ref == 'refs/heads/main'
    steps:
      - name: Promote model to production
        run: |
          python src/promote_model.py \
            --registry "image-classifier-prod" \
            --from-alias "candidate" \
            --to-alias "production"

# W&B Webhook 接收服务器示例(FastAPI)
from fastapi import FastAPI, Request
import subprocess

app = FastAPI()

@app.post("/wandb-webhook")
async def handle_wandb_webhook(request: Request):
    payload = await request.json()
    event_type = payload.get("event_type")
    artifact_alias = payload.get("artifact_alias", "")

    if event_type == "ADD_ARTIFACT_ALIAS" and "production" in artifact_alias:
        # 触发生产部署流水线
        artifact_path = payload["artifact_version_string"]
        subprocess.run([
            "kubectl", "set", "image",
            "deployment/model-server",
            f"model-server=model-image:{artifact_path}",
        ])
        # Slack 通知
        notify_slack(
            f"New production model deployed: {artifact_path}"
        )

    return {"status": "ok"}

# 以离线模式运行，之后再同步
import os
os.environ["WANDB_MODE"] = "offline"

wandb.init(project="my-project")
# ... 进行训练 ...
wandb.finish()

# 之后联网时手动同步
# 在终端中运行: wandb sync ./wandb/offline-run-*

# 使用引用型 Artifact —— 不上传实际文件，只记录路径
artifact = wandb.Artifact("large-dataset", type="dataset")
artifact.add_reference(
    "s3://my-bucket/datasets/imagenet/",  # 支持 S3、GCS、Azure Blob 等
    name="imagenet"
)
run.log_artifact(artifact)

# 按批次记录表格
BATCH_SIZE = 1000
table = wandb.Table(columns=["id", "prediction", "label"])

for i in range(0, len(predictions), BATCH_SIZE):
    batch = predictions[i:i + BATCH_SIZE]
    for pred in batch:
        table.add_data(pred["id"], pred["prediction"], pred["label"])

run.log({"predictions": table})

# 恢复已有的 sweep
wandb agent your-entity/project/existing_sweep_id

# 配置 W&B Alert —— 检测到训练异常时告警
import wandb

wandb.init(project="production-monitor")

# 训练过程中检测到异常时发送 Slack/Email 通知
if val_loss > threshold:
    wandb.alert(
        title="Training Anomaly Detected",
        text=f"Validation loss ({val_loss:.4f}) exceeded threshold ({threshold})",
        level=wandb.AlertLevel.WARN,
    )

# 检测 NaN
if torch.isnan(loss):
    wandb.alert(
        title="NaN Loss Detected",
        text="Training diverged - NaN loss detected",
        level=wandb.AlertLevel.ERROR,
    )
    wandb.finish(exit_code=1)

# 自动清理旧 Artifact 版本的脚本
import wandb

api = wandb.Api()
artifact_type = api.artifact_type("model", "your-entity/your-project")

for artifact_collection in artifact_type.collections():
    versions = artifact_collection.versions()
    # 只保留最新 5 个版本
    for version in versions[5:]:
        if "production" not in version.aliases and "staging" not in version.aliases:
            version.delete()
            print(f"Deleted: {version.name}")

import sys
import torch
import wandb

def init_wandb_with_env(project, name, config):
    """自动包含环境信息的 wandb.init 包装函数"""
    env_config = {
        "python_version": sys.version,
        "pytorch_version": torch.__version__,
        "cuda_version": torch.version.cuda,
        "cudnn_version": torch.backends.cudnn.version(),
        "gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "N/A",
        "num_gpus": torch.cuda.device_count(),
    }
    config.update(env_config)

    return wandb.init(
        project=project,
        name=name,
        config=config,
        save_code=True,  # 自动保存运行的代码
    )