""" Train Cofiber Threshold Dim20 (22K params) on full COCO 2017 train. Initialized from SVD of the pruned prototype matrix — the projection starts from the top-20 directions the pruned prototypes identified as important. Same hyperparameters as box32: batch 64, lr 1e-3, cosine + 3% warmup, 8 epochs. """ import math import os import sys import time import json import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from torch.utils.data import DataLoader, Dataset from torchvision.transforms import v2 from torchvision.ops import generalized_box_iou SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(SCRIPT_DIR, '..', '..', '..')) EUPE_REPO = os.environ.get("ARENA_BACKBONE_REPO", "/home/zootest/EUPE") EUPE_WEIGHTS = os.environ.get("ARENA_BACKBONE_WEIGHTS", "/home/zootest/weights/eupe_vitb/EUPE-ViT-B.pt") COCO_ROOT = os.environ.get("ARENA_COCO_ROOT", "/home/zootest/datasets/coco") OUTPUT_DIR = SCRIPT_DIR if EUPE_REPO not in sys.path: sys.path.insert(0, EUPE_REPO) RESOLUTION = 640 NUM_CLASSES = 80 BATCH_SIZE = 64 LR = 1e-3 WEIGHT_DECAY = 1e-4 EPOCHS = 8 GRAD_CLIP = 5.0 WARMUP_FRACTION = 0.03 COCO_CONTIG_TO_CAT = [ 1,2,3,4,5,6,7,8,9,10,11,13,14,15,16,17,18,19,20,21,22,23,24,25,27,28,31,32, 33,34,35,36,37,38,39,40,41,42,43,44,46,47,48,49,50,51,52,53,54,55,56,57,58, 59,60,61,62,63,64,65,67,70,72,73,74,75,76,77,78,79,80,81,82,84,85,86,87,88,89,90, ] COCO_CAT_TO_CONTIG = {cat: i for i, cat in enumerate(COCO_CONTIG_TO_CAT)} def letterbox(image, res): W0, H0 = image.size scale = res / max(H0, W0) new_w, new_h = int(round(W0 * scale)), int(round(H0 * scale)) resized = image.resize((new_w, new_h), Image.BILINEAR) canvas = Image.new("RGB", (res, res), (0, 0, 0)) canvas.paste(resized, (0, 0)) return canvas, scale class COCODetection(Dataset): def __init__(self, root, split="train"): img_dir = os.path.join(root, f"{split}2017") ann_file = os.path.join(root, "annotations", f"instances_{split}2017.json") with open(ann_file) as f: coco = json.load(f) self.img_dir = img_dir self.normalize = v2.Compose([ v2.ToImage(), v2.ToDtype(torch.float32, scale=True), v2.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ]) id_to_anns = {} for a in coco["annotations"]: if a["iscrowd"]: continue cat = a["category_id"] if cat not in COCO_CAT_TO_CONTIG: continue id_to_anns.setdefault(a["image_id"], []).append(a) self.items = [] id_to_info = {img["id"]: img for img in coco["images"]} for iid, anns in id_to_anns.items(): info = id_to_info[iid] boxes, labels = [], [] for a in anns: x, y, w, h = a["bbox"] if w < 1 or h < 1: continue boxes.append([x, y, x + w, y + h]) labels.append(COCO_CAT_TO_CONTIG[a["category_id"]]) if boxes: self.items.append({"file": info["file_name"], "boxes": boxes, "labels": labels}) print(f" COCO {split}: {len(self.items)} images", flush=True) def __len__(self): return len(self.items) def __getitem__(self, idx): item = self.items[idx] img = Image.open(os.path.join(self.img_dir, item["file"])).convert("RGB") canvas, scale = letterbox(img, RESOLUTION) x = self.normalize(canvas) boxes = torch.tensor(item["boxes"], dtype=torch.float32) * scale labels = torch.tensor(item["labels"], dtype=torch.long) return x, boxes, labels def collate_fn(batch): return torch.stack([b[0] for b in batch]), [b[1] for b in batch], [b[2] for b in batch] # Inline head def cofiber_decompose(f, n_scales): cofibers = [] residual = f for _ in range(n_scales - 1): omega = F.avg_pool2d(residual, 2) sigma_omega = F.interpolate(omega, size=residual.shape[2:], mode="bilinear", align_corners=False) cofibers.append(residual - sigma_omega) residual = omega cofibers.append(residual) return cofibers class CofiberThresholdDim20(nn.Module): def __init__(self, feat_dim=768, bottleneck_dim=20, num_classes=80, n_scales=3, reg_hidden=16): super().__init__() self.n_scales = n_scales self.scale_norms = nn.ModuleList([nn.LayerNorm(feat_dim) for _ in range(n_scales)]) self.project = nn.Linear(feat_dim, bottleneck_dim, bias=False) self.cls_weight = nn.Parameter(torch.randn(num_classes, bottleneck_dim) * 0.01) self.cls_bias = nn.Parameter(torch.zeros(num_classes)) self.reg_hidden_layer = nn.Linear(bottleneck_dim, reg_hidden) self.reg_act = nn.GELU() self.reg_out = nn.Linear(reg_hidden, 4) self.ctr_weight = nn.Parameter(torch.randn(1, bottleneck_dim) * 0.01) self.ctr_bias = nn.Parameter(torch.zeros(1)) self.scale_params = nn.Parameter(torch.ones(n_scales)) def forward(self, spatial): cofibers = cofiber_decompose(spatial, self.n_scales) cls_l, reg_l, ctr_l = [], [], [] for i, cof in enumerate(cofibers): B, C, H, W = cof.shape f = self.scale_norms[i](cof.permute(0, 2, 3, 1).reshape(-1, C)) z = self.project(f) cls = (z @ self.cls_weight.T + self.cls_bias).reshape(B, H, W, -1).permute(0, 3, 1, 2) reg_raw = (self.reg_out(self.reg_act(self.reg_hidden_layer(z))) * self.scale_params[i]).clamp(-10, 10) reg = torch.exp(reg_raw).reshape(B, H, W, 4).permute(0, 3, 1, 2) ctr = (z @ self.ctr_weight.T + self.ctr_bias).reshape(B, H, W, 1).permute(0, 3, 1, 2) cls_l.append(cls); reg_l.append(reg); ctr_l.append(ctr) return cls_l, reg_l, ctr_l # Inline loss (same as other scripts) def make_locations(feature_sizes, strides, device): locs = [] for (h, w), s in zip(feature_sizes, strides): ys = (torch.arange(h, device=device, dtype=torch.float32) + 0.5) * s xs = (torch.arange(w, device=device, dtype=torch.float32) + 0.5) * s gy, gx = torch.meshgrid(ys, xs, indexing="ij") locs.append(torch.stack([gx.flatten(), gy.flatten()], -1)) return locs def assign_targets(locations, boxes, labels, strides, size_ranges): cls_t, reg_t, ctr_t = [], [], [] if boxes.numel() == 0: for loc in locations: n = loc.shape[0] cls_t.append(torch.full((n,), -1, dtype=torch.long, device=loc.device)) reg_t.append(torch.zeros(n, 4, device=loc.device)) ctr_t.append(torch.zeros(n, device=loc.device)) return cls_t, reg_t, ctr_t areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) for loc, stride, sr in zip(locations, strides, size_ranges): n = loc.shape[0] l = loc[:, None, 0] - boxes[None, :, 0] t = loc[:, None, 1] - boxes[None, :, 1] r = boxes[None, :, 2] - loc[:, None, 0] b = boxes[None, :, 3] - loc[:, None, 1] ltrb = torch.stack([l, t, r, b], dim=-1) in_box = ltrb.min(dim=-1).values > 0 cx = (boxes[:, 0] + boxes[:, 2]) / 2 cy = (boxes[:, 1] + boxes[:, 3]) / 2 rad = stride * 1.5 in_center = ((loc[:, None, 0] >= cx - rad) & (loc[:, None, 0] <= cx + rad) & (loc[:, None, 1] >= cy - rad) & (loc[:, None, 1] <= cy + rad)) max_d = ltrb.max(dim=-1).values in_level = (max_d >= sr[0]) & (max_d <= sr[1]) pos = in_box & in_center & in_level a = areas[None, :].expand_as(pos).clone() a[~pos] = float("inf") matched = a.argmin(dim=-1) is_pos = a.gather(1, matched[:, None]).squeeze(1) < float("inf") ct = torch.full((n,), -1, dtype=torch.long, device=loc.device) ct[is_pos] = labels[matched[is_pos]] rt = torch.zeros(n, 4, device=loc.device) if is_pos.any(): rt[is_pos] = ltrb[torch.arange(n, device=loc.device)[is_pos], matched[is_pos]] ctrt = torch.zeros(n, device=loc.device) if is_pos.any(): lp, tp, rp, bp = rt[is_pos].unbind(-1) ctrt[is_pos] = torch.sqrt( (torch.minimum(lp, rp) / torch.maximum(lp, rp).clamp(min=1e-6)) * (torch.minimum(tp, bp) / torch.maximum(tp, bp).clamp(min=1e-6))) cls_t.append(ct); reg_t.append(rt); ctr_t.append(ctrt) return cls_t, reg_t, ctr_t def focal_loss(logits, targets, alpha=0.25, gamma=2.0): p = torch.sigmoid(logits) ce = F.binary_cross_entropy_with_logits(logits, targets, reduction="none") pt = p * targets + (1 - p) * (1 - targets) at = alpha * targets + (1 - alpha) * (1 - targets) return (at * (1 - pt) ** gamma * ce).sum() def compute_loss(cls_per, reg_per, ctr_per, locs_per, boxes_batch, labels_batch): B = cls_per[0].shape[0] device = cls_per[0].device num_classes = cls_per[0].shape[1] strides = [16, 32, 64] size_ranges = [(-1, 128), (128, 256), (256, float("inf"))] flat_cls, flat_reg, flat_ctr = [], [], [] for cl, rg, ct in zip(cls_per, reg_per, ctr_per): b, c, h, w = cl.shape flat_cls.append(cl.permute(0, 2, 3, 1).reshape(b, h * w, c)) flat_reg.append(rg.permute(0, 2, 3, 1).reshape(b, h * w, 4)) flat_ctr.append(ct.permute(0, 2, 3, 1).reshape(b, h * w)) pred_cls = torch.cat(flat_cls, 1) pred_reg = torch.cat(flat_reg, 1) pred_ctr = torch.cat(flat_ctr, 1) all_locs = torch.cat(locs_per, 0) all_ct, all_rt, all_ctt = [], [], [] for i in range(B): ct, rt, ctt = assign_targets(locs_per, boxes_batch[i], labels_batch[i], strides, size_ranges) all_ct.append(torch.cat(ct)); all_rt.append(torch.cat(rt)); all_ctt.append(torch.cat(ctt)) tgt_cls = torch.stack(all_ct) tgt_reg = torch.stack(all_rt) tgt_ctr = torch.stack(all_ctt) pos = tgt_cls >= 0 npos = max(pos.sum().item(), 1) oh = torch.zeros_like(pred_cls) pi = pos.nonzero(as_tuple=True) oh[pi[0], pi[1], tgt_cls[pos]] = 1.0 loss_cls = focal_loss(pred_cls.reshape(-1, num_classes), oh.reshape(-1, num_classes)) / npos if pos.any(): pp = pred_reg[pos]; tp = tgt_reg[pos]; pl = all_locs[None].expand(B, -1, -1)[pos] pb = torch.stack([pl[:,0]-pp[:,0], pl[:,1]-pp[:,1], pl[:,0]+pp[:,2], pl[:,1]+pp[:,3]], -1) tb = torch.stack([pl[:,0]-tp[:,0], pl[:,1]-tp[:,1], pl[:,0]+tp[:,2], pl[:,1]+tp[:,3]], -1) giou = generalized_box_iou(pb, tb) loss_reg = (1 - giou.diagonal()).sum() / npos loss_ctr = F.binary_cross_entropy_with_logits(pred_ctr[pos], tgt_ctr[pos], reduction="sum") / npos else: loss_reg = loss_ctr = torch.tensor(0.0, device=device) return loss_cls + loss_reg + loss_ctr def train(): os.makedirs(OUTPUT_DIR, exist_ok=True) print("=" * 60) print("Cofiber Threshold Dim20: 22K params, SVD-initialized, 8 epochs") print("=" * 60, flush=True) print("\n[1/4] Loading backbone...", flush=True) backbone = torch.hub.load(EUPE_REPO, "eupe_vitb16", source="local", weights=EUPE_WEIGHTS) backbone = backbone.cuda().eval() for p in backbone.parameters(): p.requires_grad = False print("\n[2/4] Building head with SVD initialization...", flush=True) head = CofiberThresholdDim20().cuda() # Initialize from SVD of pruned prototypes svd_init_path = os.path.join(SCRIPT_DIR, "svd_init.pt") if os.path.isfile(svd_init_path): svd_init = torch.load(svd_init_path, map_location="cuda") head.project.weight.data = svd_init["project"] head.cls_weight.data = svd_init["cls_weight"] print(" SVD initialization loaded", flush=True) else: print(" No SVD init found, using random", flush=True) n_params = sum(p.numel() for p in head.parameters()) print(f" {n_params:,} params", flush=True) print("\n[3/4] Loading COCO...", flush=True) train_ds = COCODetection(COCO_ROOT, "train") train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True, drop_last=True, collate_fn=collate_fn) steps_per_epoch = len(train_loader) total_steps = steps_per_epoch * EPOCHS warmup_steps = int(total_steps * WARMUP_FRACTION) print(f" {len(train_ds)} images, {steps_per_epoch} steps/epoch, {total_steps} total", flush=True) optimizer = torch.optim.AdamW(head.parameters(), lr=LR, weight_decay=WEIGHT_DECAY) def lr_lambda(step): if step < warmup_steps: return step / max(warmup_steps, 1) progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1) return 0.5 * (1.0 + math.cos(math.pi * progress)) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) strides = [16, 32, 64] H = RESOLUTION // 16 locs = make_locations([(H, H), (H//2, H//2), (H//4, H//4)], strides, torch.device("cuda")) print(f"\n[4/4] Training...", flush=True) log_file = open(os.path.join(OUTPUT_DIR, "train.log"), "a") head.train() global_step = 0 running_loss = 0.0 running_count = 0 t0 = time.time() for epoch in range(EPOCHS): print(f"\n Epoch {epoch+1}/{EPOCHS} starting (step {global_step})", flush=True) for images, boxes_b, labels_b in train_loader: if global_step >= total_steps: break images = images.cuda(non_blocking=True) boxes_b = [b.cuda(non_blocking=True) for b in boxes_b] labels_b = [l.cuda(non_blocking=True) for l in labels_b] try: with torch.no_grad(): with torch.autocast("cuda", dtype=torch.bfloat16): out = backbone.forward_features(images) patches = out["x_norm_patchtokens"].float() B, N, D = patches.shape h = w = int(N ** 0.5) spatial = patches.permute(0, 2, 1).reshape(B, D, h, w) cls_l, reg_l, ctr_l = head(spatial) loss = compute_loss(cls_l, reg_l, ctr_l, locs, boxes_b, labels_b) if torch.isnan(loss) or torch.isinf(loss): print(f" WARNING: NaN/Inf loss at step {global_step}", flush=True) optimizer.zero_grad(); global_step += 1; scheduler.step(); continue optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(head.parameters(), GRAD_CLIP) optimizer.step() scheduler.step() global_step += 1 running_loss += loss.item() running_count += 1 if global_step % 100 == 0: elapsed = time.time() - t0 avg = running_loss / max(running_count, 1) lr_now = scheduler.get_last_lr()[0] msg = f"step {global_step}/{total_steps} (epoch {epoch+1}) loss={loss.item():.4f} avg={avg:.4f} lr={lr_now:.2e} {running_count/elapsed:.1f} it/s" print(msg, flush=True) log_file.write(msg + "\n"); log_file.flush() if global_step % 1000 == 0: torch.save({"head": head.state_dict(), "global_step": global_step}, os.path.join(OUTPUT_DIR, "checkpoint.pth")) print(f" Checkpoint saved at step {global_step}", flush=True) except Exception as e: import traceback print(f"\n ERROR at step {global_step}: {e}", flush=True) traceback.print_exc() if "out of memory" in str(e): torch.cuda.empty_cache(); optimizer.zero_grad(); global_step += 1; scheduler.step(); continue raise print(f" Epoch {epoch+1}/{EPOCHS} complete (step {global_step})", flush=True) final_path = os.path.join(OUTPUT_DIR, "cofiber_threshold_dim20_coco_8ep_22k.pth") torch.save(head.state_dict(), final_path) print(f"\nSaved: {final_path}") print(f"Training complete: {total_steps} steps, {(time.time()-t0)/3600:.1f} hours", flush=True) log_file.close() if __name__ == "__main__": train()