evaluate.py

"""Evaluate a trained checkpoint on the MNIST test set.

    python evaluate.py
    python evaluate.py --checkpoint checkpoints/mnist.npz --errors-png errors.png

Prints overall accuracy, a per-class breakdown, and the confusion matrix, and
can dump the images the network got wrong — which is the fastest way to tell a
model that is genuinely weak from a test set that contains genuinely ambiguous
handwriting.
"""

import argparse
import os

import numpy as np

import augment as augmentation
import data
from nn import load_checkpoint

HERE = os.path.dirname(os.path.abspath(__file__))


def parse_args():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--checkpoint", default=None,
                   help="default: most recently trained checkpoint")
    p.add_argument("--split", choices=["test", "val", "train"], default="test")
    p.add_argument("--errors-png", nargs="?", const="auto", default=None,
                   help="save a grid of the most confident mistakes")
    p.add_argument("--tta", type=int, default=0, metavar="N",
                   help="test-time augmentation: average over N warped views "
                        "plus the clean one")
    p.add_argument("--tta-rotation", type=float, default=5.0, help="degrees")
    p.add_argument("--tta-translate", type=float, default=1.0, help="pixels")
    p.add_argument("--tta-scale", type=float, default=0.05, help="fraction")
    p.add_argument("--tta-elastic", action="store_true",
                   help="include elastic distortion in the TTA views")
    p.add_argument("--seed", type=int, default=0)
    return p.parse_args()


def predict_probs(model, x, batch_size=1000):
    """Softmax probabilities for a whole split, in eval mode."""
    logits = np.concatenate([model.forward(x[i:i + batch_size], training=False)
                             for i in range(0, len(x), batch_size)])
    z = logits - logits.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)


def tta_probs(model, x, warp, rng, n_views, batch_size=1000):
    """Average predictions over the clean image plus `n_views` warped copies.

    Two choices worth stating.  The clean image is always included and carries
    the same weight as any other view: the test set is clean, so throwing away
    the one view drawn from the actual test distribution to average only warps
    would be strange.

    And it averages *probabilities*, not logits.  Averaging logits is a
    geometric mean once you push it through the softmax, which lets a single
    view that is confidently wrong drag the whole ensemble with it — the
    arithmetic mean of probabilities lets the other views outvote it.
    """
    total = predict_probs(model, x, batch_size)
    for _ in range(n_views):
        total = total + predict_probs(model, warp(x, rng), batch_size)
    return total / (n_views + 1)


def confusion_matrix(y_true, y_pred, num_classes):
    cm = np.zeros((num_classes, num_classes), dtype=np.int64)
    np.add.at(cm, (y_true, y_pred), 1)
    return cm


def print_report(y_true, y_pred, num_classes):
    cm = confusion_matrix(y_true, y_pred, num_classes)
    support = cm.sum(axis=1)
    predicted = cm.sum(axis=0)
    tp = np.diag(cm)

    with np.errstate(divide="ignore", invalid="ignore"):
        recall = np.where(support > 0, tp / np.maximum(support, 1), 0.0)
        precision = np.where(predicted > 0, tp / np.maximum(predicted, 1), 0.0)
        f1 = np.where(precision + recall > 0,
                      2 * precision * recall / np.maximum(precision + recall, 1e-12),
                      0.0)

    print("\nper class")
    print(f"  {'digit':>5}  {'precision':>9}  {'recall':>7}  {'f1':>6}  "
          f"{'support':>7}  {'errors':>6}")
    for c in range(num_classes):
        print(f"  {c:>5}  {precision[c] * 100:8.2f}%  {recall[c] * 100:6.2f}%  "
              f"{f1[c]:6.3f}  {support[c]:>7}  {support[c] - tp[c]:>6}")

    print("\nconfusion matrix (rows = true, cols = predicted, . = 0)")
    print("        " + "".join(f"{c:>6}" for c in range(num_classes)))
    for c in range(num_classes):
        cells = "".join(f"{v:>6}" if v else f"{'.':>6}" for v in cm[c])
        print(f"  true {c}{cells}")

    # The single most-confusable pair is usually more informative than the
    # whole matrix: it says what the network actually struggles to tell apart.
    off = cm.copy()
    np.fill_diagonal(off, 0)
    i, j = np.unravel_index(off.argmax(), off.shape)
    if off[i, j]:
        print(f"\nworst confusion: true {i} predicted as {j}, {off[i, j]} times")


def save_error_grid(path, images, y_true, y_pred, conf, mean, std, limit=36):
    """Grid of the mistakes the network was most confident about."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    order = np.argsort(-conf)[:limit]
    if len(order) == 0:
        print("no errors to plot")
        return

    cols = 6
    rows = int(np.ceil(len(order) / cols))
    fig, axes = plt.subplots(rows, cols, figsize=(cols * 1.6, rows * 1.75))
    for ax in np.atleast_1d(axes).ravel():
        ax.axis("off")

    for ax, idx in zip(np.atleast_1d(axes).ravel(), order):
        img = images[idx] * std + mean          # undo normalisation for display
        ax.imshow(img.reshape(28, 28), cmap="gray_r", vmin=0, vmax=1)
        ax.set_title(f"{y_true[idx]}→{y_pred[idx]}  {conf[idx] * 100:.0f}%",
                     fontsize=8)

    fig.suptitle("most confident mistakes (true → predicted)", fontsize=11)
    fig.tight_layout()
    fig.savefig(path, dpi=140)
    print(f"error grid: {path}")


def newest_checkpoint():
    """Whatever was trained last — the usual thing you want to look at."""
    folder = os.path.join(HERE, "checkpoints")
    found = [os.path.join(folder, f) for f in os.listdir(folder)
             if f.endswith(".npz")] if os.path.isdir(folder) else []
    if not found:
        raise SystemExit("no checkpoints found — run train.py first")
    return max(found, key=os.path.getmtime)


def main():
    args = parse_args()
    checkpoint = args.checkpoint or newest_checkpoint()
    if not os.path.exists(checkpoint):
        raise SystemExit(f"no checkpoint at {checkpoint} — run train.py first")

    model, config, extra = load_checkpoint(checkpoint)
    # The checkpoint knows its own architecture, so the data layout follows
    # from it rather than from a flag the caller has to remember.  It also
    # knows how the data was split: rebuilding the split from *this* script's
    # defaults would put training images into `--split val` whenever training
    # used a different --seed or --val-size, and would re-standardise every
    # split with constants the network was never trained under.
    split_seed = int(extra["seed"]) if "seed" in extra else args.seed
    split_val = int(extra["val_size"]) if "val_size" in extra else 5000
    d = data.load(seed=split_seed, val_size=split_val,
                  flatten=config.get("arch", "mlp") == "mlp")
    x, y = d[f"x_{args.split}"], d[f"y_{args.split}"]

    # Older checkpoints predate seed/val_size but do carry mean/std, so the
    # reconstruction can still be checked against what training actually used.
    if "mean" in extra and abs(extra["mean"] - d["mean"]) > 1e-6:
        print(f"WARNING: this split normalises to mean {d['mean']:.7f}, but "
              f"the checkpoint was trained at {extra['mean']:.7f} — the split "
              f"could not be reconstructed, so these numbers are not exact.")

    probs = predict_probs(model, x)
    y_pred = probs.argmax(axis=1)

    acc = float((y_pred == y).mean())
    print(f"checkpoint: {checkpoint}")
    print(f"trained to epoch {extra.get('epoch', '?')} "
          f"(val {extra.get('val_acc', float('nan')) * 100:.2f}%)")
    print(model)
    print(f"\n{args.split} accuracy {acc * 100:.2f}%  "
          f"({int((y_pred != y).sum())} errors / {len(y)})")

    if args.tta:
        warp = augmentation.build(d["image_shape"], fill=d["fill"],
                                  affine=True, elastic=args.tta_elastic,
                                  rotation=args.tta_rotation,
                                  translate=args.tta_translate,
                                  scale=args.tta_scale)
        probs = tta_probs(model, x, warp, np.random.default_rng(args.seed),
                          args.tta)
        y_pred = probs.argmax(axis=1)
        tta_acc = float((y_pred == y).mean())
        # Print the plain number above and this one here, always both: TTA that
        # is not compared against the un-augmented baseline is unfalsifiable.
        print(f"\ntta ({args.tta} views + clean) {warp}")
        print(f"{args.split} accuracy {tta_acc * 100:.2f}%  "
              f"({int((y_pred != y).sum())} errors / {len(y)})  "
              f"[{(tta_acc - acc) * 100:+.2f} pts]")

    print_report(y, y_pred, config["num_classes"])

    if args.errors_png:
        path = (os.path.splitext(checkpoint)[0] + f"_{args.split}_errors.png"
                if args.errors_png == "auto" else args.errors_png)
        wrong = np.where(y_pred != y)[0]
        save_error_grid(path, x[wrong], y[wrong], y_pred[wrong],
                        probs[wrong].max(axis=1), d["mean"], d["std"])


if __name__ == "__main__":
    main()