train_fashionmnist_wandb.py

# Reference: https://chatgpt.com/c/69b72e10-7f8c-832e-9845-6a532587cb13
# ============================================================
# FashionMNIST training script with Weights & Biases (W&B) logging.
#
# W&B key concepts used in this script:
#   wandb.init()         – start a new experiment run and upload config
#   wandb.config         – read hyperparameters synced from init()
#   wandb.log()          – stream metrics (loss, accuracy, …) to the dashboard
#   wandb.watch()        – auto-log model gradients and parameter histograms
#   wandb.Image/Table    – log rich media (images, structured tables)
#   wandb.plot.*         – log built-in interactive charts (e.g. confusion matrix)
#   wandb.run.summary    – store final / best scalar values for a run
#   wandb.save()         – upload files (checkpoints) as run artifacts
#   wandb.finish()       – mark the run as complete and flush any remaining data
# ============================================================

import os
import random

import numpy as np
import torch
import wandb

from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


CLASS_NAMES = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]


# ============================================================
# CONFIG
# ============================================================
def get_config():
    # This dict is passed directly to wandb.init(config=...), which uploads
    # every key-value pair so you can compare runs in the W&B dashboard.
    config = {
        "project_name": "fashionmnist-pytorch-demo",
        "group_name": "model_comparison",
        "model_name": "cnn",   # options: "mlp" or "cnn"
        "data_dir": "data",
        "checkpoint_dir": "checkpoints",
        "batch_size": 128,
        "learning_rate": 1e-3,  # options: 1e-2, 1e-3, 1e-4
        "epochs": 10,
        "hidden_dim_1": 256,
        "hidden_dim_2": 128,
        "cnn_channels_1": 32,
        "cnn_channels_2": 64,
        "num_classes": 10,
        "seed": 42,
    }

    if config["model_name"] == "mlp":
        config["run_name"] = (
            f"mlp-lr{config['learning_rate']}-"
            f"bs{config['batch_size']}-"
            f"h1{config['hidden_dim_1']}-"
            f"h2{config['hidden_dim_2']}-"
            f"seed{config['seed']}"
        )
    elif config["model_name"] == "cnn":
        config["run_name"] = (
            f"cnn-lr{config['learning_rate']}-"
            f"bs{config['batch_size']}-"
            f"c1{config['cnn_channels_1']}-"
            f"c2{config['cnn_channels_2']}-"
            f"seed{config['seed']}"
        )
    else:
        raise ValueError(f"Unsupported model_name: {config['model_name']}")

    return config


# ============================================================
# SETUP
# ============================================================
def seed_everything(seed: int = 42):
    """
    Set random seeds for reproducibility.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)


def get_device():
    """
    Choose device in priority order:
    CUDA > MPS > CPU
    """
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")
    return device


# ============================================================
# DATA
# ============================================================
def build_dataloaders(config):
    transform = transforms.ToTensor()

    train_dataset = datasets.FashionMNIST(
        root=config["data_dir"],
        train=True,
        download=True,
        transform=transform,
    )

    test_dataset = datasets.FashionMNIST(
        root=config["data_dir"],
        train=False,
        download=True,
        transform=transform,
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=config["batch_size"],
        shuffle=True,
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=config["batch_size"],
        shuffle=False,
    )

    return train_loader, test_loader


# ============================================================
# MODEL
# ============================================================
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10, c1=32, c2=64):
        super().__init__()

        self.features = nn.Sequential(
            nn.Conv2d(1, c1, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),   # 28x28 -> 14x14

            nn.Conv2d(c1, c2, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),   # 14x14 -> 7x7
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(c2 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x


def build_model(config):
    if config["model_name"] == "mlp":
        model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, config["hidden_dim_1"]),
            nn.ReLU(),
            nn.Linear(config["hidden_dim_1"], config["hidden_dim_2"]),
            nn.ReLU(),
            nn.Linear(config["hidden_dim_2"], config["num_classes"]),
        )
    elif config["model_name"] == "cnn":
        model = SimpleCNN(
            num_classes=config["num_classes"],
            c1=config["cnn_channels_1"],
            c2=config["cnn_channels_2"],
        )
    else:
        raise ValueError(f"Unsupported model_name: {config['model_name']}")

    return model


# ============================================================
# TRAIN
# ============================================================
def train_one_epoch(model, dataloader, loss_fn, optimizer, device, epoch):
    model.train()

    total_loss = 0.0
    total_samples = 0

    for batch_idx, (images, labels) in enumerate(dataloader):
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        outputs = model(images)
        loss = loss_fn(outputs, labels)

        loss.backward()
        optimizer.step()

        batch_size = images.size(0)
        total_loss += loss.item() * batch_size
        total_samples += batch_size

        if batch_idx % 100 == 0:
            # wandb.log() streams a dict of metrics to the W&B dashboard.
            # Keys with a "/" prefix (e.g. "train/") create sub-sections in the UI.
            # Each call auto-increments an internal step counter used on the x-axis.
            wandb.log(
                {
                    "train/batch_loss": loss.item(),
                    "train/epoch": epoch,
                    "train/batch_idx": batch_idx,
                }
            )

    return total_loss / total_samples


# ============================================================
# EVAL
# ============================================================
@torch.no_grad()
def evaluate(model, dataloader, loss_fn, device):
    model.eval()

    total_loss = 0.0
    total_correct = 0
    total_samples = 0

    for images, labels in dataloader:
        images = images.to(device)
        labels = labels.to(device)

        outputs = model(images)
        loss = loss_fn(outputs, labels)

        batch_size = images.size(0)
        total_loss += loss.item() * batch_size
        total_samples += batch_size

        predictions = outputs.argmax(dim=1)
        total_correct += (predictions == labels).sum().item()

    avg_loss = total_loss / total_samples
    accuracy = total_correct / total_samples

    return avg_loss, accuracy


@torch.no_grad()
def log_vision_predictions(model, dataloader, device, max_samples=24, max_misclassified=24):
    model.eval()

    all_true = []
    all_pred = []

    sample_rows = []
    misclassified_rows = []

    for images, labels in dataloader:
        images = images.to(device)
        labels = labels.to(device)

        outputs = model(images)
        predictions = outputs.argmax(dim=1)

        all_true.extend(labels.cpu().tolist())
        all_pred.extend(predictions.cpu().tolist())

        images_cpu = images.cpu()
        labels_cpu = labels.cpu()
        predictions_cpu = predictions.cpu()

        for i in range(images_cpu.size(0)):
            true_idx = labels_cpu[i].item()
            pred_idx = predictions_cpu[i].item()

            # wandb.Image() wraps a tensor/PIL image for rich media logging.
            # The caption is overlaid in the W&B media browser.
            image_for_wandb = wandb.Image(
                images_cpu[i],
                caption=f"true={CLASS_NAMES[true_idx]}, pred={CLASS_NAMES[pred_idx]}"
            )

            if len(sample_rows) < max_samples:
                sample_rows.append([
                    image_for_wandb,
                    CLASS_NAMES[true_idx],
                    CLASS_NAMES[pred_idx],
                    true_idx == pred_idx,
                ])

            if true_idx != pred_idx and len(misclassified_rows) < max_misclassified:
                misclassified_rows.append([
                    image_for_wandb,
                    CLASS_NAMES[true_idx],
                    CLASS_NAMES[pred_idx],
                ])

        if len(sample_rows) >= max_samples and len(misclassified_rows) >= max_misclassified:
            break

    # wandb.Table() creates an interactive table in the W&B UI.
    # You can filter, sort, and view embedded media (images) directly in the table.
    sample_table = wandb.Table(
        columns=["image", "true_label", "pred_label", "correct"],
        data=sample_rows,
    )

    misclassified_table = wandb.Table(
        columns=["image", "true_label", "pred_label"],
        data=misclassified_rows,
    )

    # wandb.plot.confusion_matrix() is a built-in chart — no custom code needed.
    # It renders an interactive heatmap in the W&B dashboard.
    confusion = wandb.plot.confusion_matrix(
        preds=all_pred,
        y_true=all_true,
        class_names=CLASS_NAMES,
        title="Test Confusion Matrix",
    )

    # Log all media in one call so they share the same step in the timeline.
    wandb.log({
        "eval/sample_predictions": sample_table,
        "eval/misclassified_examples": misclassified_table,
        "eval/confusion_matrix": confusion,
    })


def main():
    # ----------------------------
    # config / setup
    # ----------------------------
    config = get_config()
    seed_everything(config["seed"])
    device = get_device()

    print(f"Using device: {device}")
    print(f"Group name: {config['group_name']}")
    print(f"Run name: {config['run_name']}")

    # wandb.init() creates a new run inside the given project.
    #   project   – groups all runs of this experiment in the W&B UI.
    #   name      – human-readable label for this specific run.
    #   group     – lets you visually cluster related runs (e.g. mlp vs cnn).
    #   config    – uploads the hyperparameter dict; accessible as wandb.config.
    wandb.init(
        project=config["project_name"],
        name=config["run_name"],
        group=config["group_name"],
        config=config,
    )
    # Add any values not known at init time (here: which device was selected).
    # allow_val_change=True is required when updating a key that already exists.
    wandb.config.update({"device": str(device)}, allow_val_change=True)

    checkpoint_dir = os.path.join(config["checkpoint_dir"], config["model_name"])
    os.makedirs(checkpoint_dir, exist_ok=True)

    # ----------------------------
    # data
    # ----------------------------
    train_loader, test_loader = build_dataloaders(config)

    # ----------------------------
    # model
    # ----------------------------
    model = build_model(config).to(device)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=config["learning_rate"])

    # wandb.watch() hooks into PyTorch to automatically record:
    #   log="all"     – both gradients AND parameter tensors each step.
    #   log_freq=100  – only record every 100 optimizer steps (reduces overhead).
    # Remove this call if you only care about scalar metrics.
    wandb.watch(model, loss_fn, log="all", log_freq=100)

    # ----------------------------
    # train / eval loop
    # ----------------------------
    best_accuracy = 0.0

    for epoch in range(1, config["epochs"] + 1):
        train_loss = train_one_epoch(
            model=model,
            dataloader=train_loader,
            loss_fn=loss_fn,
            optimizer=optimizer,
            device=device,
            epoch=epoch,
        )

        val_loss, val_accuracy = evaluate(
            model=model,
            dataloader=test_loader,
            loss_fn=loss_fn,
            device=device,
        )

        print(
            f"Epoch [{epoch}/{config['epochs']}] "
            f"train_loss={train_loss:.4f} "
            f"val_loss={val_loss:.4f} "
            f"val_accuracy={val_accuracy:.4f}"
        )

        # Log epoch-level metrics alongside the step-level batch metrics above.
        # Using the same wandb.log() call keeps all metrics on a shared x-axis.
        wandb.log(
            {
                "epoch": epoch,
                "train/epoch_loss": train_loss,
                "val/loss": val_loss,
                "val/accuracy": val_accuracy,
            }
        )

        if val_accuracy > best_accuracy:
            best_accuracy = val_accuracy
            best_model_path = os.path.join(checkpoint_dir, "best_model.pt")
            torch.save(model.state_dict(), best_model_path)
            # wandb.run.summary stores single "final" values for a run.
            # Unlike wandb.log(), summary values are not time-series — they appear
            # in the runs table so you can sort/filter all runs by best metric.
            wandb.run.summary["best_val_accuracy"] = best_accuracy
            print(f"Saved best model to: {best_model_path}")

    # ----------------------------
    # final save
    # ----------------------------
    final_model_path = os.path.join(checkpoint_dir, "final_model.pt")
    torch.save(model.state_dict(), final_model_path)
    print(f"Saved final model to: {final_model_path}")

    log_vision_predictions(
        model=model,
        dataloader=test_loader,
        device=device,
        max_samples=24,
        max_misclassified=24,
    )

    # wandb.save() uploads matched files as artifacts attached to this run.
    # They will be available for download from the W&B run page.
    wandb.save(os.path.join(checkpoint_dir, "*.pt"))

    # wandb.finish() finalises the run: flushes remaining logs, marks it
    # as "finished" in the dashboard, and closes the background sync thread.
    # Always call this at the end of a script (or use wandb.init() as a context manager).
    wandb.finish()


if __name__ == "__main__":
    main()