Weights & Biases: Experiment Tracking, Sweeps & Reports

Getting Started: wandb.init() and Logging

Install with pip install wandb, authenticate with wandb login, then add three lines to any training script:

import wandb
import torch

# Initialize a run — creates a unique run in your project
run = wandb.init(
    project="image-classifier",
    name="resnet50-lr-0.001",
    config={
        "architecture": "ResNet50",
        "learning_rate": 1e-3,
        "batch_size": 64,
        "epochs": 30,
        "optimizer": "AdamW",
    }
)

for epoch in range(config.epochs):
    train_loss, val_loss, val_acc = train_epoch(model, ...)

    # Log metrics — visible in real time on wandb.ai
    wandb.log({
        "epoch": epoch,
        "train/loss": train_loss,
        "val/loss": val_loss,
        "val/accuracy": val_acc,
        "lr": scheduler.get_last_lr()[0],
    })

    # Log images, confusion matrices, histograms
    wandb.log({"predictions": wandb.Image(sample_img, caption="epoch " + str(epoch))})

# Save model as artifact
artifact = wandb.Artifact("resnet50-classifier", type="model")
artifact.add_file("model.pth")
run.log_artifact(artifact)
wandb.finish()

W&B Sweeps: Automated Hyperparameter Search

Define a sweep config, then launch agents that automatically sample and run experiments:

import wandb

# Define the search space and strategy
sweep_config = {
    "method": "bayes",        # bayes | random | grid
    "metric": {"name": "val/accuracy", "goal": "maximize"},
    "parameters": {
        "learning_rate": {"distribution": "log_uniform_values", "min": 1e-5, "max": 1e-1},
        "batch_size": {"values": [16, 32, 64, 128]},
        "dropout": {"distribution": "uniform", "min": 0.1, "max": 0.5},
        "optimizer": {"values": ["adam", "adamw", "sgd"]},
    },
    "early_terminate": {"type": "hyperband", "min_iter": 3},
}

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

def train_sweep():
    with wandb.init() as run:
        cfg = run.config
        model = build_model(dropout=cfg.dropout)
        optimizer = get_optimizer(cfg.optimizer, cfg.learning_rate)
        # ... training loop
        wandb.log({"val/accuracy": final_acc})

# Launch 10 parallel agents
wandb.agent(sweep_id, function=train_sweep, count=50)

HuggingFace Transformers Integration

Zero-config logging with HuggingFace Trainer — just set report_to="wandb":

from transformers import TrainingArguments, Trainer
import wandb

wandb.init(project="llm-finetuning", name="llama3-8b-qlora")

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    learning_rate=2e-4,
    report_to="wandb",                # ← enable W&B logging
    logging_steps=10,
    evaluation_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=500,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
)
# W&B auto-logs: loss curves, learning rate, GPU util, eval metrics, model checkpoints

W&B vs MLflow vs Comet vs Neptune