backend/evaluate.py

"""
Evaluation script for the Deep Ensemble.

Computes:
- Accuracy on CIFAR-10 test set
- Expected Calibration Error (ECE)
- OOD detection AUROC (CIFAR-10 vs CIFAR-100 using predictive entropy)

Usage:
    python evaluate.py
    python evaluate.py --checkpoint-dir ./checkpoints --ensemble-size 5
"""

import argparse
import os

import numpy as np
import torch
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as T

from model import get_model
from utils import (
    ensemble_predict, mean_prediction, predictive_entropy,
    mutual_information, expected_calibration_error, ood_auroc,
    CIFAR10_CLASSES,
)


def load_ensemble(checkpoint_dir, ensemble_size, device, num_classes=10):
    """Load all ensemble member models from checkpoints."""
    models = []
    for i in range(ensemble_size):
        path = os.path.join(checkpoint_dir, f"model_{i}.pt")
        if not os.path.exists(path):
            print(f"Warning: Checkpoint not found: {path}")
            continue

        model = get_model(num_classes=num_classes, depth=16, widen_factor=2)
        checkpoint = torch.load(path, map_location=device, weights_only=True)
        model.load_state_dict(checkpoint["model_state_dict"])
        model.to(device)
        model.eval()
        models.append(model)
        print(f"  Loaded model_{i} (test_acc: {checkpoint.get('test_acc', '?'):.2f}%, "
              f"seed: {checkpoint.get('seed', '?')})")

    print(f"Loaded {len(models)}/{ensemble_size} ensemble members")
    return models


def get_test_loader(dataset_class, data_dir="./data", batch_size=128, num_workers=4):
    """Create a test dataloader with standard normalization."""
    transform = T.Compose([
        T.ToTensor(),
        T.Normalize((0.4914, 0.4822, 0.4465),
                     (0.2470, 0.2435, 0.2616)),
    ])
    dataset = dataset_class(
        root=data_dir, train=False, download=True, transform=transform)
    return DataLoader(
        dataset, batch_size=batch_size, shuffle=False,
        num_workers=num_workers, pin_memory=True)


def evaluate_accuracy_and_calibration(models, loader, device):
    """Compute ensemble accuracy and ECE on a dataset."""
    all_probs = []
    all_targets = []

    for inputs, targets in loader:
        probs = ensemble_predict(models, inputs, device)  # (M, B, K)
        all_probs.append(probs)
        all_targets.append(targets.numpy())

    all_probs = np.concatenate(all_probs, axis=1)    # (M, N, K)
    all_targets = np.concatenate(all_targets, axis=0)  # (N,)

    mean_p = mean_prediction(all_probs)                # (N, K)
    predictions = np.argmax(mean_p, axis=1)            # (N,)
    confidences = np.max(mean_p, axis=1)               # (N,)
    accuracies_bin = (predictions == all_targets).astype(float)

    accuracy = 100.0 * accuracies_bin.mean()
    ece = expected_calibration_error(confidences, accuracies_bin)

    return accuracy, ece, all_probs


def compute_ood_scores(models, id_loader, ood_loader, device):
    """Compute uncertainty scores for ID and OOD data."""
    id_probs_list = []
    ood_probs_list = []

    print("Computing ID (CIFAR-10) uncertainty scores...")
    for inputs, _ in id_loader:
        probs = ensemble_predict(models, inputs, device)
        id_probs_list.append(probs)

    print("Computing OOD (CIFAR-100) uncertainty scores...")
    for inputs, _ in ood_loader:
        probs = ensemble_predict(models, inputs, device)
        ood_probs_list.append(probs)

    id_probs = np.concatenate(id_probs_list, axis=1)    # (M, N_id, K)
    ood_probs = np.concatenate(ood_probs_list, axis=1)  # (M, N_ood, K)

    # Predictive entropy as OOD score
    id_entropy = predictive_entropy(mean_prediction(id_probs))
    ood_entropy = predictive_entropy(mean_prediction(ood_probs))

    # Mutual information as OOD score
    id_mi = mutual_information(id_probs)
    ood_mi = mutual_information(ood_probs)

    return {
        "id_entropy": id_entropy,
        "ood_entropy": ood_entropy,
        "id_mi": id_mi,
        "ood_mi": ood_mi,
    }


def main():
    parser = argparse.ArgumentParser(description="Evaluate Deep Ensemble")
    parser.add_argument("--checkpoint-dir", type=str, default="./checkpoints")
    parser.add_argument("--ensemble-size", type=int, default=5)
    parser.add_argument("--data-dir", type=str, default="./data")
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument("--gpu", type=int, default=0)
    args = parser.parse_args()

    # Device
    if torch.cuda.is_available():
        device = torch.device(f"cuda:{args.gpu}")
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")
    print(f"Device: {device}")

    # Load ensemble
    print("\nLoading ensemble...")
    models = load_ensemble(args.checkpoint_dir, args.ensemble_size, device)
    if not models:
        print("ERROR: No models loaded. Train models first with train.py")
        return

    # Evaluate on CIFAR-10 (in-distribution)
    print("\n" + "=" * 60)
    print("IN-DISTRIBUTION EVALUATION (CIFAR-10)")
    print("=" * 60)
    cifar10_loader = get_test_loader(
        torchvision.datasets.CIFAR10, args.data_dir,
        args.batch_size, args.num_workers)
    accuracy, ece, _ = evaluate_accuracy_and_calibration(
        models, cifar10_loader, device)
    print(f"  Accuracy: {accuracy:.2f}%")
    print(f"  ECE:      {ece:.4f}")

    # OOD detection (CIFAR-10 vs CIFAR-100)
    print("\n" + "=" * 60)
    print("OOD DETECTION (CIFAR-10 vs CIFAR-100)")
    print("=" * 60)
    cifar100_loader = get_test_loader(
        torchvision.datasets.CIFAR100, args.data_dir,
        args.batch_size, args.num_workers)
    scores = compute_ood_scores(models, cifar10_loader, cifar100_loader, device)

    auroc_entropy = ood_auroc(scores["id_entropy"], scores["ood_entropy"])
    auroc_mi = ood_auroc(scores["id_mi"], scores["ood_mi"])
    print(f"  AUROC (Predictive Entropy): {auroc_entropy:.4f}")
    print(f"  AUROC (Mutual Information): {auroc_mi:.4f}")

    # Summary statistics
    print("\n" + "=" * 60)
    print("UNCERTAINTY STATISTICS")
    print("=" * 60)
    print(f"  ID Entropy  — mean: {scores['id_entropy'].mean():.4f}, "
          f"std: {scores['id_entropy'].std():.4f}")
    print(f"  OOD Entropy — mean: {scores['ood_entropy'].mean():.4f}, "
          f"std: {scores['ood_entropy'].std():.4f}")
    print(f"  ID MI       — mean: {scores['id_mi'].mean():.4f}, "
          f"std: {scores['id_mi'].std():.4f}")
    print(f"  OOD MI      — mean: {scores['ood_mi'].mean():.4f}, "
          f"std: {scores['ood_mi'].std():.4f}")

    print("\n✓ Evaluation complete.")


if __name__ == "__main__":
    main()