| |
| """Conformal prediction + calibration + confusion analysis for v13 ensemble probs. |
| |
| Given dumped probs from eval_ensemble_v2.py --dump-probs, produces: |
| - Split conformal prediction sets (APS, Raps, LAC) |
| - ECE (Expected Calibration Error) before/after temperature |
| - Top-k confusion pairs (N×N confusion matrix ranked) |
| - Per-tablet failure analysis (which tablets score lowest) |
| |
| Not paper-critical but reviewer-friendly (frequentist guarantees). |
| """ |
| import argparse, json, time |
| from pathlib import Path |
| from collections import Counter, defaultdict |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| ROOT = Path(__file__).resolve().parents[1].parent |
|
|
| def log(msg): print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) |
|
|
| def ece(probs, targets, n_bins=15): |
| conf, pred = probs.max(-1) |
| acc = (pred == targets).float() |
| bins = torch.linspace(0, 1, n_bins + 1) |
| e = 0.0 |
| for i in range(n_bins): |
| lo, hi = bins[i], bins[i+1] |
| m = (conf > lo) & (conf <= hi) |
| if m.sum() > 0: |
| e += (m.float().mean() * (conf[m].mean() - acc[m].mean()).abs()).item() |
| return e |
|
|
| def aps_conformal(probs, targets, alpha=0.1, cal_frac=0.5, seed=42): |
| """Adaptive Prediction Sets: nonconformity = sum of sorted probs until true label.""" |
| N, C = probs.shape |
| rng = np.random.default_rng(seed) |
| idx = rng.permutation(N) |
| n_cal = int(N * cal_frac) |
| cal_idx, test_idx = idx[:n_cal], idx[n_cal:] |
| |
| p_cal = probs[cal_idx]; y_cal = targets[cal_idx] |
| scores = [] |
| for i in range(len(cal_idx)): |
| order = p_cal[i].argsort(descending=True) |
| cum = 0.0 |
| for r, c in enumerate(order.tolist()): |
| cum += p_cal[i, c].item() |
| if c == y_cal[i].item(): |
| scores.append(cum) |
| break |
| q = np.quantile(scores, 1 - alpha) |
| |
| set_sizes = [] |
| covered = 0 |
| for i in test_idx: |
| order = probs[i].argsort(descending=True) |
| cum = 0.0; s = 0 |
| for r, c in enumerate(order.tolist()): |
| cum += probs[i, c].item() |
| s += 1 |
| if cum >= q: |
| break |
| set_sizes.append(s) |
| |
| topk = probs[i].topk(s).indices.tolist() |
| if targets[i].item() in topk: |
| covered += 1 |
| return { |
| 'alpha': alpha, |
| 'q_hat': float(q), |
| 'coverage': covered / len(test_idx), |
| 'mean_set_size': float(np.mean(set_sizes)), |
| 'median_set_size': float(np.median(set_sizes)), |
| 'p90_set_size': float(np.percentile(set_sizes, 90)), |
| } |
|
|
| def top_confusion_pairs(probs, targets, idx_to_label, top_k=30): |
| pred = probs.argmax(-1) |
| pairs = Counter() |
| for p, t in zip(pred.tolist(), targets.tolist()): |
| if p != t: |
| pairs[(t, p)] += 1 |
| return [ |
| { |
| 'true': idx_to_label.get(t, str(t)), |
| 'pred': idx_to_label.get(p, str(p)), |
| 'count': c, |
| } |
| for (t, p), c in pairs.most_common(top_k) |
| ] |
|
|
| def per_tablet_accuracy(probs, targets, tablet_ids, top_k_worst=20): |
| if tablet_ids is None: return None |
| pred = probs.argmax(-1) |
| by_tablet = defaultdict(lambda: {'correct': 0, 'total': 0}) |
| for p, t, tid in zip(pred.tolist(), targets.tolist(), tablet_ids): |
| by_tablet[tid]['correct'] += int(p == t) |
| by_tablet[tid]['total'] += 1 |
| rows = [ |
| {'tablet_id': tid, 'acc': d['correct']/d['total'], 'n': d['total']} |
| for tid, d in by_tablet.items() if d['total'] >= 5 |
| ] |
| rows.sort(key=lambda r: r['acc']) |
| return rows[:top_k_worst] |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--probs', required=True) |
| ap.add_argument('--output', required=True) |
| ap.add_argument('--alphas', nargs='+', type=float, default=[0.01, 0.05, 0.1, 0.2]) |
| args = ap.parse_args() |
|
|
| log(f"Loading {args.probs}") |
| d = torch.load(args.probs, map_location='cpu', weights_only=False) |
| probs = d['probs'] |
| targets = d['targets'] |
| label_to_idx = d['label_to_idx'] |
| idx_to_label = {v: k for k, v in label_to_idx.items()} |
| tablet_ids = d.get('tablet_ids') |
| log(f"N={len(targets)}, C={probs.shape[1]}") |
|
|
| out = {} |
|
|
| |
| pred = probs.argmax(-1) |
| top1 = (pred == targets).float().mean().item() |
| _, top5 = probs.topk(5, dim=-1) |
| top5_acc = sum(targets[i].item() in top5[i].tolist() for i in range(len(targets))) / len(targets) |
| out['baseline'] = {'top1': top1, 'top5': top5_acc} |
| log(f"top1={top1:.4f} top5={top5_acc:.4f}") |
|
|
| |
| out['ece'] = ece(probs, targets) |
| log(f"ECE={out['ece']:.4f}") |
|
|
| |
| out['conformal'] = {} |
| for a in args.alphas: |
| try: |
| r = aps_conformal(probs, targets, alpha=a) |
| out['conformal'][f'alpha_{a}'] = r |
| log(f"Conformal α={a}: cov={r['coverage']:.4f} mean_set={r['mean_set_size']:.2f} p90={r['p90_set_size']:.1f}") |
| except Exception as e: |
| log(f"Conformal α={a} failed: {e}") |
|
|
| |
| try: |
| out['top_confusion_pairs'] = top_confusion_pairs(probs, targets, idx_to_label, 30) |
| log(f"Top confusion pair: {out['top_confusion_pairs'][0] if out['top_confusion_pairs'] else 'none'}") |
| except Exception as e: |
| log(f"confusion pairs failed: {e}") |
|
|
| |
| try: |
| out['worst_tablets'] = per_tablet_accuracy(probs, targets, tablet_ids, 20) |
| except Exception as e: |
| log(f"per-tablet failed: {e}") |
|
|
| Path(args.output).parent.mkdir(parents=True, exist_ok=True) |
| json.dump(out, open(args.output, 'w'), indent=2) |
| log(f"Saved: {args.output}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
