""" Magnitude pruning of the trained cofiber threshold prototypes. For each INT8 weight, try reducing |w| by 1. If detection precision holds within 95% of unpruned, accept the reduction. Batched evaluation on the cached COCO diagnostic images. Usage: python prune.py --passes 5 """ import argparse import os import sys import time import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', '..')) def load_head_and_data(): """Load the trained head and diagnostic validation data.""" from heads.cofiber_threshold.head import CofiberThreshold # Load trained weights ckpt_path = os.path.join(os.path.dirname(__file__), 'cofiber_threshold_coco_8ep_70k.pth') head = CofiberThreshold() head.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=False)) head = head.cuda().eval() # Load cached COCO validation features cache_dir = os.environ.get('ARENA_CACHE_DIR', '/mnt/d/JacobProject/outputs/arena_cache') val_path = os.path.join(cache_dir, 'domains', 'coco', 'val.pt') val_data = torch.load(val_path, map_location='cpu', weights_only=False) val_items = val_data['items'] return head, val_items def evaluate_detection(head, val_items): """Run detection on val items, return (tp, fp, n_gt).""" from utils.decode import make_locations, decode_fcos # Get locations sample = val_items[0]['spatial'].unsqueeze(0).float().cuda() locs = head.get_locs(sample) total_tp, total_fp, total_gt = 0, 0, 0 with torch.no_grad(): for item in val_items: sp = item['spatial'].unsqueeze(0).float().cuda() cls_l, reg_l, ctr_l = head(sp) dets = decode_fcos(cls_l, reg_l, ctr_l, locs, score_thresh=0.3) det_boxes = dets[0]['boxes'] det_labels = dets[0]['labels'] gt_boxes = item['boxes'].cuda() gt_labels = item['labels'].cuda() n_gt = len(gt_labels) total_gt += n_gt if len(det_boxes) > 0 and n_gt > 0: x1 = torch.maximum(det_boxes[:, None, 0], gt_boxes[None, :, 0]) y1 = torch.maximum(det_boxes[:, None, 1], gt_boxes[None, :, 1]) x2 = torch.minimum(det_boxes[:, None, 2], gt_boxes[None, :, 2]) y2 = torch.minimum(det_boxes[:, None, 3], gt_boxes[None, :, 3]) inter = (x2 - x1).clamp(min=0) * (y2 - y1).clamp(min=0) d_area = (det_boxes[:, 2] - det_boxes[:, 0]) * (det_boxes[:, 3] - det_boxes[:, 1]) g_area = (gt_boxes[:, 2] - gt_boxes[:, 0]) * (gt_boxes[:, 3] - gt_boxes[:, 1]) iou = inter / (d_area[:, None] + g_area[None, :] - inter).clamp(min=1e-6) matched = set() for di in range(len(det_boxes)): best_iou, best_gi = iou[di].max(0) gi = best_gi.item() if best_iou.item() >= 0.5 and gi not in matched and det_labels[di] == gt_labels[gi]: total_tp += 1 matched.add(gi) else: total_fp += 1 else: total_fp += len(det_boxes) precision = total_tp / max(total_tp + total_fp, 1) return precision, total_tp, total_fp, total_gt def prune_weights(head, val_items, passes=5, threshold=0.95): """Iteratively reduce weight magnitudes while maintaining detection precision.""" base_prec, base_tp, base_fp, base_gt = evaluate_detection(head, val_items) min_prec = base_prec * threshold print(f"Base precision: {base_prec:.4f} (tp={base_tp}, fp={base_fp}, gt={base_gt})") print(f"Minimum acceptable: {min_prec:.4f}") total_reductions = 0 initial_magnitude = sum(p.abs().sum().item() for p in head.parameters()) initial_nonzero = sum((p != 0).sum().item() for p in head.parameters()) for pass_num in range(passes): pass_reductions = 0 t0 = time.time() for name, param in head.named_parameters(): if param.numel() < 10: # skip tiny params (biases, scales) continue flat = param.data.flatten() for i in range(len(flat)): val = flat[i].item() if val == 0: continue # Try reducing magnitude by a small step new_val = val - 0.001 if val > 0 else val + 0.001 flat[i] = new_val param.data = flat.reshape(param.shape) prec, _, _, _ = evaluate_detection(head, val_items) if prec >= min_prec: pass_reductions += 1 else: # Revert flat[i] = val param.data = flat.reshape(param.shape) current_mag = sum(p.abs().sum().item() for p in head.parameters()) current_nonzero = sum((p != 0).sum().item() for p in head.parameters()) reduction_pct = 100 * (1 - current_mag / initial_magnitude) elapsed = time.time() - t0 total_reductions += pass_reductions print(f"Pass {pass_num+1}: {pass_reductions} reductions, " f"magnitude -{reduction_pct:.2f}%, " f"nonzero {current_nonzero}/{initial_nonzero}, " f"time {elapsed:.0f}s") if pass_reductions == 0: print("No reductions possible. Stopping.") break final_prec, tp, fp, gt = evaluate_detection(head, val_items) final_mag = sum(p.abs().sum().item() for p in head.parameters()) final_nonzero = sum((p != 0).sum().item() for p in head.parameters()) print(f"\nPruning complete:") print(f" Precision: {base_prec:.4f} -> {final_prec:.4f}") print(f" Magnitude: {initial_magnitude:.0f} -> {final_mag:.0f} (-{100*(1-final_mag/initial_magnitude):.1f}%)") print(f" Nonzero: {initial_nonzero} -> {final_nonzero} (-{initial_nonzero - final_nonzero})") print(f" Total reductions: {total_reductions}") return head def main(): parser = argparse.ArgumentParser() parser.add_argument('--passes', type=int, default=3) parser.add_argument('--threshold', type=float, default=0.95) args = parser.parse_args() print("=" * 60) print("Cofiber Threshold Pruning") print("=" * 60) head, val_items = load_head_and_data() head = prune_weights(head, val_items, passes=args.passes, threshold=args.threshold) # Save pruned weights out_path = os.path.join(os.path.dirname(__file__), 'cofiber_threshold_coco_8ep_70k_pruned.pth') torch.save(head.state_dict(), out_path) print(f"\nSaved: {out_path}") if __name__ == '__main__': main()