trained/linear_70k/prune_batched.py

"""
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()