Weights & Biases 初探：给 MNIST CNN 加个实验追踪

931 字

5 分钟

Weights & Biases 初探：给 MNIST CNN 加个实验追踪

2025-12-04

机器学习

MachineLearning

/

wandb

/

PyTorch

/

MNIST

起因#

昨天偶然接触到一个使用 wandb 的项目。打开官网看了看介绍，发现这东西还挺有意思的。

Weights & Biases（简称 wandb）是一个机器学习实验追踪平台。它能自动记录你训练过程中的各种指标、超参数、模型权重，甚至是 git commit 信息。官方宣传的卖点是”几行代码集成，永久保存实验记录”。

听起来很美好。于是今天决定亲自试一试。

实验设置#

让 Codex 写了一个简单的 MNIST CNN，然后用 wandb 进行可视化追踪。整个过程确实挺顺滑的——wandb.init() 初始化项目，wandb.log() 记录指标，wandb.watch() 追踪梯度。setup 一次，后面基本不用管。

主要代码如下：

1
def train(config):
2
    device = get_device()
3
    print(f"Using device: {device}")
4
    wandb.init(project=config.project, config=config, mode=config.wandb_mode)
5
    model = SimpleCNN().to(device)
6
    wandb.watch(model, log="gradients", log_freq=100)
7

8
    train_loader, val_loader = get_dataloaders(config.batch_size)
9
    criterion = nn.CrossEntropyLoss()
10
    optimizer = optim.Adam(model.parameters(), lr=config.lr)
11

12
    for epoch in range(1, config.epochs + 1):
13
        model.train()
14
        running_loss = 0.0
15
        correct = 0
16
        total = 0
17

18
        for batch_idx, (images, labels) in enumerate(train_loader, start=1):
19
            images, labels = images.to(device), labels.to(device)
20
            optimizer.zero_grad()
21
            outputs = model(images)
22
            loss = criterion(outputs, labels)
23
            loss.backward()
24
            optimizer.step()
25

26
            running_loss += loss.item() * images.size(0)
27
            _, predicted = torch.max(outputs, 1)
28
            total += labels.size(0)
29
            correct += (predicted == labels).sum().item()
30

31
        train_loss = running_loss / total
32
        train_acc = correct / total
33

34
        val_loss, val_acc = evaluate(model, val_loader, device)
35

36
        log_data = {
37
            "epoch": epoch,
38
            "train/loss": train_loss,
39
            "train/accuracy": train_acc,
40
            "val/loss": val_loss,
41
            "val/accuracy": val_acc,
42
        }
43

44
        if config.log_samples:
45
            sample_images, sample_labels = next(iter(val_loader))
46
            sample_images = sample_images.to(device)
47
            with torch.no_grad():
48
                preds = model(sample_images).argmax(dim=1)
49
            log_data["val/sample_predictions"] = log_sample_predictions(
50
                sample_images[:16].cpu(), sample_labels[:16], preds[:16].cpu()
51
            )
52

53
        wandb.log(log_data)
54
        print(
55
            f"Epoch {epoch}/{config.epochs} "
56
            f"- train_loss: {train_loss:.4f}, train_acc: {train_acc:.4f} "
57
            f"- val_loss: {val_loss:.4f}, val_acc: {val_acc:.4f}"
58
        )
59

60
    wandb.finish()
61
    return model