""" Batched INT8 magnitude pruning for the cofiber threshold detector. Quantizes the trained prototypes to INT8, then tests weight reductions in parallel batches on GPU. Accepts reductions that maintain detection precision within 95% of the unpruned baseline. Adapted from the 8bit-threshold-computer prune_weights.py methodology. """ 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__), '..', '..', '..')) from heads.cofiber_threshold.head import CofiberThreshold, cofiber_decompose from utils.decode import make_locations, decode_fcos CACHE_DIR = os.environ.get('ARENA_CACHE_DIR', '/mnt/d/JacobProject/outputs/arena_cache') def load_val_data(): val_path = os.path.join(CACHE_DIR, 'domains', 'coco', 'val.pt') val_data = torch.load(val_path, map_location='cpu', weights_only=False) return val_data['items'] def evaluate_precision(head, val_items, locs): """Fast precision eval on cached val items.""" 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) return total_tp / max(total_tp + total_fp, 1), total_tp, total_fp, total_gt def main(): print("=" * 60) print("Batched INT8 Pruning: Cofiber Threshold") print("=" * 60) # Load print("\nLoading head and val data...") 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() val_items = load_val_data() sample = val_items[0]['spatial'].unsqueeze(0).float().cuda() locs = head.get_locs(sample) # Baseline print("\nBaseline evaluation...") base_prec, base_tp, base_fp, base_gt = evaluate_precision(head, val_items, locs) print(f" Precision: {base_prec:.4f} (tp={base_tp}, fp={base_fp}, gt={base_gt})") min_prec = base_prec * 0.95 print(f" Threshold: {min_prec:.4f} (95% of baseline)") # Quantize prototypes to INT8 scale proto = head.prototypes.data scale = proto.abs().max().item() / 127.0 print(f"\n Prototype scale: {scale:.6f}") print(f" Prototype range: [{proto.min().item():.4f}, {proto.max().item():.4f}]") # Pruning: try zeroing each weight one at a time initial_nonzero = (proto.abs() > scale).sum().item() initial_magnitude = proto.abs().sum().item() print(f" Initial nonzero (above 1 INT8 unit): {initial_nonzero}") print(f" Initial magnitude: {initial_magnitude:.1f}") total_pruned = 0 passes = 3 for pass_num in range(passes): t0 = time.time() pass_pruned = 0 # Sort weights by magnitude — try smallest first flat = proto.flatten() sorted_idx = flat.abs().argsort() for rank, idx in enumerate(sorted_idx): val = flat[idx].item() if abs(val) < scale * 0.5: continue # already effectively zero in INT8 # Try reducing magnitude by one INT8 step new_val = val - scale if val > 0 else val + scale if abs(new_val) < scale * 0.5: new_val = 0.0 flat[idx] = new_val head.prototypes.data = flat.reshape(proto.shape) prec, _, _, _ = evaluate_precision(head, val_items, locs) if prec >= min_prec: pass_pruned += 1 if pass_pruned % 100 == 0: elapsed = time.time() - t0 current_mag = flat.abs().sum().item() print(f" Pass {pass_num+1}: {pass_pruned} pruned, " f"mag {current_mag:.0f} (-{100*(1-current_mag/initial_magnitude):.1f}%), " f"prec {prec:.4f}, {elapsed:.0f}s", flush=True) else: # Revert flat[idx] = val head.prototypes.data = flat.reshape(proto.shape) if rank > 5000: break # cap per pass total_pruned += pass_pruned current_mag = head.prototypes.data.abs().sum().item() current_zeros = (head.prototypes.data.abs() < scale * 0.5).sum().item() elapsed = time.time() - t0 print(f" Pass {pass_num+1}: {pass_pruned} reductions, " f"zeros {current_zeros}/{proto.numel()}, " f"mag -{100*(1-current_mag/initial_magnitude):.1f}%, " f"time {elapsed:.0f}s", flush=True) if pass_pruned == 0: break # Final eval final_prec, tp, fp, gt = evaluate_precision(head, val_items, locs) final_zeros = (head.prototypes.data.abs() < scale * 0.5).sum().item() final_nonzero = proto.numel() - final_zeros final_mag = head.prototypes.data.abs().sum().item() print(f"\n{'='*60}") print(f"PRUNING RESULTS") print(f"{'='*60}") print(f" Precision: {base_prec:.4f} -> {final_prec:.4f}") print(f" Nonzero prototypes: {proto.numel()} -> {final_nonzero} ({final_zeros} zeroed)") print(f" Magnitude: {initial_magnitude:.0f} -> {final_mag:.0f} (-{100*(1-final_mag/initial_magnitude):.1f}%)") print(f" Total reductions: {total_pruned}") # Save 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" Saved: {out_path}") if __name__ == '__main__': main()