reproduction-bundle/src/evaluate_source_classifier.py

#!/usr/bin/env python3
"""Evaluate the pinned LOGDIFF CMNIST classifiers on the complete test split.

The upstream trainer reports accuracy at one random diffusion timestep per
validation example.  This script adds a deterministic audit at fixed timesteps,
keeps the digit and colour heads separate, and records checkpoint provenance.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import random
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
import torchvision
from diffusers import DDPMScheduler


ROOT = Path(__file__).resolve().parents[2]
UPSTREAM = ROOT / "vendor" / "LogDiff" / "logdiff"
sys.path.insert(0, str(UPSTREAM))

from cs_classifier.models import MultiLabelClassifier, MultiheadClassifier  # noqa: E402
from datasets.coloredMNIST import ColoredMNIST  # noqa: E402


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def device_from_arg(name: str) -> torch.device:
    if name != "auto":
        return torch.device(name)
    if torch.backends.mps.is_available():
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def load_checkpoint(path: Path, kind: str, device: torch.device) -> torch.nn.Module:
    backbone = torchvision.models.resnet18(weights=None)
    if kind == "composition":
        model = MultiheadClassifier(backbone, [10, 10])
    else:
        model = MultiLabelClassifier(backbone, [10, 10])
    payload = torch.load(path, map_location="cpu", weights_only=False)
    model.load_state_dict(payload["state_dict"], strict=True)
    return model.eval().to(device)


@torch.inference_mode()
def evaluate(
    model: torch.nn.Module,
    loader: torch.utils.data.DataLoader,
    kind: str,
    timesteps: list[int],
    device: torch.device,
    seed: int,
) -> list[dict[str, object]]:
    scheduler = DDPMScheduler(
        num_train_timesteps=1000,
        clip_sample=True,
        prediction_type="epsilon",
        beta_schedule="linear",
    )
    alpha = scheduler.alphas_cumprod.float()
    rows: list[dict[str, object]] = []
    eval_steps: list[int | None] = timesteps if kind == "composition" else [None]

    for timestep in eval_steps:
        correct = torch.zeros(2, dtype=torch.long)
        nll_sum = torch.zeros(2, dtype=torch.float64)
        total = 0
        generator = torch.Generator(device="cpu").manual_seed(seed + (timestep or 0))
        for batch in loader:
            clean = batch["X"]
            labels = batch["label"]
            if timestep is None:
                model_input = clean.to(device)
                logits = model(model_input)
            else:
                noise = torch.randn(clean.shape, generator=generator)
                a = alpha[timestep]
                noisy = a.sqrt() * clean + (1.0 - a).sqrt() * noise
                model_input = noisy.to(device)
                t = torch.full((clean.shape[0],), timestep, device=device, dtype=torch.long)
                logits = model(model_input, t)

            labels_device = labels.to(device)
            for head in range(2):
                correct[head] += (logits[head].argmax(1) == labels_device[:, head]).sum().cpu()
                nll_sum[head] += F.cross_entropy(
                    logits[head], labels_device[:, head], reduction="sum"
                ).cpu().double()
            total += clean.shape[0]

        rows.append(
            {
                "kind": kind,
                "timestep": "clean" if timestep is None else timestep,
                "examples": total,
                "digit_accuracy": float(correct[0] / total),
                "colour_accuracy": float(correct[1] / total),
                "digit_nll": float(nll_sum[0] / total),
                "colour_nll": float(nll_sum[1] / total),
            }
        )
    return rows


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--kind", choices=["composition", "judge"], required=True)
    parser.add_argument("--data-root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, default=ROOT / "outputs" / "classifier-audit")
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--timesteps", type=int, nargs="+", default=[0, 50, 100, 250, 500, 750, 999])
    args = parser.parse_args()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    device = device_from_arg(args.device)
    dataset = ColoredMNIST(args.data_root, train=False, download=True, support="uniform")
    loader = torch.utils.data.DataLoader(
        dataset, batch_size=args.batch_size, shuffle=False, num_workers=0
    )
    model = load_checkpoint(args.checkpoint, args.kind, device)
    rows = evaluate(model, loader, args.kind, args.timesteps, device, args.seed)

    args.output_dir.mkdir(parents=True, exist_ok=True)
    csv_path = args.output_dir / f"{args.kind}_fixed_timestep_metrics.csv"
    with csv_path.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    report = {
        "paper": "OAM1jJsMGp",
        "upstream_commit": "94ef35bafd4b4239e9832d8295128c09e8fc1472",
        "kind": args.kind,
        "checkpoint": args.checkpoint.name,
        "checkpoint_sha256": sha256(args.checkpoint),
        "device": str(device),
        "test_split_examples": len(dataset),
        "metrics": rows,
    }
    json_path = args.output_dir / f"{args.kind}_fixed_timestep_metrics.json"
    json_path.write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()