| |
| """Mega ensemble — union all trained heads by probabilistic averaging. |
| |
| Inputs: list of pt files, each containing {probs, targets, top1}. |
| Optimizes weights via coordinate descent on val top1. |
| """ |
| import json, argparse, itertools |
| from pathlib import Path |
| import torch |
| import torch.nn.functional as F |
| import numpy as np |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--probs', nargs='+', required=True, help='name:path pt files') |
| ap.add_argument('--optimize', action='store_true') |
| ap.add_argument('--output', required=True) |
| args = ap.parse_args() |
|
|
| names, probs, targets = [], [], None |
| own_acc = [] |
| for spec in args.probs: |
| if ':' in spec: |
| nm, pth = spec.split(':', 1) |
| else: |
| nm = Path(spec).stem; pth = spec |
| if not Path(pth).exists(): |
| print(f"skip missing {pth}"); continue |
| d = torch.load(pth, map_location='cpu', weights_only=False) |
| p = d.get('probs') |
| if p is None: p = d.get('refined_probs') |
| t = d.get('targets') |
| if p is None or t is None: |
| print(f"skip {pth}: no probs/targets"); continue |
| if targets is None: |
| targets = t |
| elif not torch.equal(t, targets): |
| print(f"{nm}: target mismatch (size {t.size(0)} vs {targets.size(0)}); skip") |
| continue |
| probs.append(p) |
| names.append(nm) |
| acc = (p.argmax(-1) == targets).float().mean().item() |
| own_acc.append(acc) |
| print(f"{nm}: top1={acc:.4f}") |
|
|
| if not probs: |
| print("No probs to ensemble"); return |
|
|
| stacked = torch.stack(probs).double() |
| M = stacked.size(0) |
|
|
| def topk(w, k=1): |
| w = torch.tensor(w, dtype=stacked.dtype).view(-1, 1, 1) |
| p = (stacked * w).sum(0) |
| if k == 1: return (p.argmax(-1) == targets).float().mean().item() |
| _, tk = p.topk(k, -1) |
| return sum(targets[i].item() in tk[i].tolist() for i in range(len(targets))) / len(targets) |
|
|
| |
| own = np.array(own_acc, dtype=np.float64) |
| w_uni = np.ones(M) / M |
| w_own = own / own.sum() |
| best_w = w_own if topk(w_own) > topk(w_uni) else w_uni |
| best_s = max(topk(w_own), topk(w_uni)) |
| print(f"Seed top1: uniform={topk(w_uni):.4f}, own-acc={topk(w_own):.4f}") |
|
|
| if args.optimize: |
| step = 0.1 |
| for it in range(100): |
| improved = False |
| for i, j in itertools.combinations(range(M), 2): |
| for d in (step, -step): |
| w = best_w.copy() |
| w[i] += d; w[j] -= d |
| if w.min() < 0: continue |
| w = w / w.sum() |
| a = topk(w) |
| if a > best_s + 1e-6: |
| best_w, best_s = w, a; improved = True |
| if not improved: step *= 0.5 |
| if step < 1e-4: break |
| print(f"Optimized top1: {best_s:.4f}") |
|
|
| |
| w_t = torch.tensor(best_w, dtype=stacked.dtype).view(-1, 1, 1) |
| final_probs = (stacked * w_t).sum(0) |
| top1 = (final_probs.argmax(-1) == targets).float().mean().item() |
| _, top5 = final_probs.topk(5, -1) |
| top5_acc = sum(targets[i].item() in top5[i].tolist() for i in range(len(targets))) / len(targets) |
|
|
| |
| max_p = final_probs.max(-1).values |
| pred = final_probs.argmax(-1) |
| sel = {} |
| for thr in (0.5, 0.6, 0.7, 0.8, 0.9): |
| keep = max_p >= thr |
| cov = keep.float().mean().item() |
| acc = (pred[keep] == targets[keep]).float().mean().item() if keep.sum() else 0.0 |
| sel[str(thr)] = {'selective_acc': acc, 'coverage': cov} |
|
|
| print(f"=== MEGA ENSEMBLE ===") |
| print(f"Top-1: {top1:.4f}, Top-5: {top5_acc:.4f}") |
| for t, m in sel.items(): |
| print(f" τ={t}: sel_acc={m['selective_acc']:.4f} cov={m['coverage']:.4f}") |
|
|
| out = { |
| 'weights': dict(zip(names, best_w.tolist())), |
| 'top1': top1, 'top5': top5_acc, |
| 'selective_metrics': sel, |
| 'per_model_top1': dict(zip(names, own_acc)), |
| 'n_val': len(targets), |
| } |
| Path(args.output).parent.mkdir(parents=True, exist_ok=True) |
| json.dump(out, open(args.output, 'w'), indent=2) |
| torch.save({'probs': final_probs, 'targets': targets}, |
| args.output.replace('.json', '_probs.pt')) |
| print(f"Saved → {args.output}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
