File size: 4,061 Bytes
0716d4a 938f692 0716d4a 938f692 0716d4a 938f692 0716d4a 938f692 0716d4a 938f692 0716d4a 938f692 0716d4a 938f692 0716d4a 938f692 0716d4a 938f692 0716d4a 938f692 0716d4a | 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 | """Single-head ablation sweep over the 144 (block, head) pairs.
For each head, zero the columns of its block's attention output projection,
score the calibration pool with the Stage 0 classifier, and record the F1 delta
and the L2 deviation of the 40 classifier-relevant output dims. Heads are then
ranked by ascending F1 drop, so the head whose removal helps most sorts first,
and the cumulative curve prunes that prefix.
Writes head_importance.json and pruning_curve.json.
"""
import argparse
import json
import sys
import time
from pathlib import Path
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # repo root, for `common`
from common import (BACKBONE, N_BLOCKS, N_HEADS, device, f1_at, heads_masked, # noqa: E402
load_pool, score_pool, write_artifact)
from common.models import load_backbone # noqa: E402
from common.pools import CALIB1000, by_name # noqa: E402
HERE = Path(__file__).resolve().parent
CLASSIFIER = HERE.parent / 'stage_0' / 'classifier.json'
CURVE_K = [1, 5, 10, 15, 20, 30, 40, 50, 60, 80, 100, 120, 144]
RANKING = 'ascending F1_drop; the smallest drop is the most prunable head'
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument('--backbone', default=BACKBONE)
ap.add_argument('--pool', default=CALIB1000.name)
args = ap.parse_args()
dev = device()
c = json.loads(CLASSIFIER.read_text())
pos = torch.tensor(c['pos_dims'], dtype=torch.long, device=dev)
neg = torch.tensor(c['neg_dims'], dtype=torch.long, device=dev)
target_dims = torch.cat([pos, neg]).unique()
thr = float(c['threshold'])
print(f'[init] |target_dims|={len(target_dims)} threshold={thr:.3f}', flush=True)
backbone = load_backbone(args.backbone).to(dev)
pool = by_name(args.pool)
print(f'[pool] {pool.name}, preloading', flush=True)
loaded = load_pool(pool, dev, preload=True)
print(f' person rate {loaded.positive_rate:.3f}', flush=True)
base_scores, base_targets = score_pool(backbone, loaded, pos, neg, target_dims)
base = f1_at(base_scores, loaded.labels, thr)
print(f'[baseline] F1={base.f1:.4f} P={base.precision:.4f} R={base.recall:.4f}',
flush=True)
results = []
t0 = time.time()
for b in range(N_BLOCKS):
for h in range(N_HEADS):
with heads_masked(backbone, [(b, h)]):
scores, targets = score_pool(backbone, loaded, pos, neg, target_dims)
m = f1_at(scores, loaded.labels, thr)
l2 = (targets - base_targets).pow(2).sum(dim=1).sqrt().mean().item()
results.append({'block': b, 'head': h, 'F1': m.f1, 'precision': m.precision,
'recall': m.recall, 'F1_drop': base.f1 - m.f1, 'target_L2': l2})
print(f' B{b:>2}H{h:>2} F1={m.f1:.4f} drop={base.f1 - m.f1:+.4f} '
f'L2={l2:.3f} {time.time() - t0:.1f}s', flush=True)
ranked = sorted(results, key=lambda r: r['F1_drop'])
curve = []
for k in CURVE_K:
with heads_masked(backbone, [(r['block'], r['head']) for r in ranked[:k]]):
scores, _ = score_pool(backbone, loaded, pos, neg)
m = f1_at(scores, loaded.labels, thr)
curve.append({'heads_pruned': k, 'F1': m.f1, 'F1_drop': base.f1 - m.f1,
'precision': m.precision, 'recall': m.recall})
print(f' K={k:>3} F1={m.f1:.4f} drop={base.f1 - m.f1:+.4f}', flush=True)
stamp = dict(generator='stage_2/ablate.py', classifier=CLASSIFIER,
pool_info=loaded.provenance())
write_artifact(HERE / 'head_importance.json', {
'ranking': RANKING,
'baseline': base.asdict(),
'per_head': results,
'ranked_most_prunable_first': [(r['block'], r['head'], r['F1_drop'])
for r in ranked],
}, **stamp)
write_artifact(HERE / 'pruning_curve.json',
{'baseline': base.asdict(), 'curve': curve}, **stamp)
print('[done]', flush=True)
if __name__ == '__main__':
main()
|