File size: 4,413 Bytes
f211247 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | #!/usr/bin/env python3
"""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, N, C)
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)
# Seed: own-acc proportional
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}")
# Final
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)
# Selective metrics
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()
|