| """Eval the deep evolved head on COCO val.""" |
| import os, sys, time, torch |
| import torch.nn.functional as F |
|
|
| SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, SCRIPT_DIR) |
|
|
| COCO_ROOT = os.environ["ARENA_COCO_ROOT"] |
| VAL_CACHE = os.environ["ARENA_VAL_CACHE"] |
| DEVICE = "cuda" |
|
|
| from train_evolved_deep import EvolvedDeepHead, cofiber_decompose, make_locations |
|
|
| |
| import json |
| with open(os.path.join(SCRIPT_DIR, "circuit", "evolved_extreme.json")) as f: |
| evolved = json.load(f) |
| dims = sorted(list(set([r for r in evolved if r["K"] == 100][0]["genome"]))) |
|
|
| |
| head = EvolvedDeepHead(dims, hidden=128, n_layers=10).to(DEVICE) |
| ckpt = os.path.join(SCRIPT_DIR, "heads", "cofiber_threshold", "evolved_deep", |
| "evolved_deep_128h_10l_20ep.pth") |
| head.load_state_dict(torch.load(ckpt, map_location=DEVICE, weights_only=False)) |
| head.eval() |
| n_params = sum(p.numel() for p in head.parameters()) |
| print(f"Loaded: {n_params:,} params") |
|
|
| |
| val = torch.load(VAL_CACHE, map_location="cpu", weights_only=False) |
| from pycocotools.coco import COCO |
| from pycocotools.cocoeval import COCOeval |
| coco_gt = COCO(os.path.join(COCO_ROOT, "annotations", "instances_val2017.json")) |
| cat_ids = sorted(coco_gt.getCatIds()) |
| idx_to_cat = {i: c for i, c in enumerate(cat_ids)} |
|
|
| strides = [16, 32, 64]; H = 640 // 16 |
| sizes = [(H,H),(H//2,H//2),(H//4,H//4)] |
| all_locs = torch.cat(make_locations(sizes, strides, torch.device(DEVICE))) |
|
|
| all_results = [] |
| t0 = time.time() |
| with torch.no_grad(): |
| for idx in range(len(val)): |
| item = val[idx] |
| spatial = item["spatial"].unsqueeze(0).float().to(DEVICE) |
| img_id = int(item["img_id"]); scale = item["scale"] |
|
|
| cls_l, reg_l, ctr_l = head(spatial) |
|
|
| cls_s = torch.cat([c.permute(0,2,3,1).reshape(-1, 80) for c in cls_l]).sigmoid() |
| reg_s = torch.cat([r.permute(0,2,3,1).reshape(-1, 4) for r in reg_l]) |
| ctr_s = torch.cat([c.permute(0,2,3,1).reshape(-1) for c in ctr_l]).sigmoid() |
|
|
| scores = cls_s * ctr_s.unsqueeze(1) |
| max_s, max_c = scores.max(1) |
| topk = min(100, max_s.shape[0]) |
| top_s, top_i = max_s.topk(topk) |
| tc = max_c[top_i]; tr = reg_s[top_i]; tl = all_locs[top_i] |
| x1=(tl[:,0]-tr[:,0])/scale; y1=(tl[:,1]-tr[:,1])/scale |
| x2=(tl[:,0]+tr[:,2])/scale; y2=(tl[:,1]+tr[:,3])/scale |
| w=(x2-x1).clamp(min=0); h=(y2-y1).clamp(min=0) |
| for i in range(topk): |
| s = top_s[i].item() |
| if s < 0.01: continue |
| all_results.append({"image_id": img_id, "category_id": idx_to_cat[tc[i].item()], |
| "bbox": [x1[i].item(), y1[i].item(), w[i].item(), h[i].item()], |
| "score": s}) |
| if (idx+1) % 1000 == 0: |
| print(f" {idx+1}/{len(val)} ({time.time()-t0:.0f}s)", flush=True) |
|
|
| print(f"\n{len(all_results)} detections") |
| if all_results: |
| coco_dt = coco_gt.loadRes(all_results) |
| coco_eval = COCOeval(coco_gt, coco_dt, "bbox") |
| coco_eval.params.imgIds = sorted(coco_gt.getImgIds())[:len(val)] |
| coco_eval.evaluate(); coco_eval.accumulate(); coco_eval.summarize() |
| mAP = coco_eval.stats[0]; mAP50 = coco_eval.stats[1]; mAP75 = coco_eval.stats[2] |
| print(f"\nDeep Evolved Head: {n_params:,} params") |
| print(f" mAP@[.5:.95]={mAP:.4f} mAP@.50={mAP50:.4f} mAP@.75={mAP75:.4f}") |
|
|
