""" Train Cofiber Threshold (S) to convergence on full COCO 2017 train. Matches the shipping FCOS hyperparameters exactly: batch 64, lr 1e-3, cosine + 3% warmup, 8 epochs, AdamW, grad clip 5.0 The only variable is the head architecture. """ import math import os import sys import time import json import numpy as np 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 sys.path.insert(0, os.path.dirname(__file__)) 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", "/mnt/d/JacobProject/datasets/llava_instruct/coco") OUTPUT_DIR = os.environ.get("ARENA_OUTPUT_DIR", "/mnt/d/detection-heads/outputs/cofiber_threshold_full") 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)} # --------------------------------------------------------------------------- # Dataset # --------------------------------------------------------------------------- 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") 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): images = torch.stack([b[0] for b in batch]) boxes = [b[1] for b in batch] labels = [b[2] for b in batch] return images, boxes, labels # --------------------------------------------------------------------------- # Cofiber Threshold head (inlined to avoid import chain issues in WSL) # --------------------------------------------------------------------------- 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 CofiberThreshold(nn.Module): def __init__(self, feat_dim=768, num_classes=NUM_CLASSES, n_scales=3): super().__init__() self.n_scales = n_scales self.scale_norms = nn.ModuleList([nn.LayerNorm(feat_dim) for _ in range(n_scales)]) self.prototypes = nn.Parameter(torch.randn(num_classes, feat_dim) * 0.01) self.proto_bias = nn.Parameter(torch.zeros(num_classes)) self.reg_weight = nn.Parameter(torch.randn(4, feat_dim) * 0.01) self.reg_bias = nn.Parameter(torch.zeros(4)) self.ctr_weight = nn.Parameter(torch.randn(1, feat_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)) cls = (f @ self.prototypes.T + self.proto_bias).reshape(B, H, W, -1).permute(0, 3, 1, 2) raw = ((f @ self.reg_weight.T + self.reg_bias) * self.scale_params[i]).clamp(-10, 10) reg = torch.exp(raw).reshape(B, H, W, 4).permute(0, 3, 1, 2) ctr = (f @ 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 # --------------------------------------------------------------------------- # FCOS loss (inlined) # --------------------------------------------------------------------------- from torchvision.ops import generalized_box_iou, nms 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, strides, size_ranges): B = cls_per[0].shape[0] device = cls_per[0].device num_classes = cls_per[0].shape[1] 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 # --------------------------------------------------------------------------- # Training # --------------------------------------------------------------------------- def train(): os.makedirs(OUTPUT_DIR, exist_ok=True) print("=" * 60) print("Full COCO training: Cofiber Threshold (S)") print("=" * 60) # Backbone print("\n[1/4] Loading backbone...") 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 # Head print("\n[2/4] Building Cofiber Threshold head...") head = CofiberThreshold().cuda() n_params = sum(p.numel() for p in head.parameters()) print(f" {n_params:,} params ({n_params/1e3:.1f}K)") # Dataset print("\n[3/4] Loading COCO...") 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, {warmup_steps} warmup") # Optimizer + schedule 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) # Precompute locations strides = [16, 32, 64] size_ranges = [(-1, 128), (128, 256), (256, float("inf"))] H = RESOLUTION // 16 feature_sizes = [(H, H), (H // 2, H // 2), (H // 4, H // 4)] locs = make_locations(feature_sizes, strides, torch.device("cuda")) # Training loop print(f"\n[4/4] Training...") log_path = os.path.join(OUTPUT_DIR, "train.log") log_file = open(log_path, "a") global_step = 0 ckpt_path = os.path.join(OUTPUT_DIR, "checkpoint.pth") if os.path.isfile(ckpt_path): print(f" Resuming from {ckpt_path}") ckpt = torch.load(ckpt_path, map_location="cuda") head.load_state_dict(ckpt["head"]) global_step = ckpt["global_step"] for _ in range(global_step): scheduler.step() print(f" Resumed at step {global_step}") head.train() running_loss = 0.0 running_count = 0 t0 = time.time() for epoch in range(EPOCHS): if global_step >= (epoch + 1) * steps_per_epoch: continue 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: # Backbone forward 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) # Head forward cls_l, reg_l, ctr_l = head(spatial) # Loss loss = compute_loss(cls_l, reg_l, ctr_l, locs, boxes_b, labels_b, strides, size_ranges) if torch.isnan(loss) or torch.isinf(loss): 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] vram = torch.cuda.max_memory_allocated() / 1024**3 msg = (f"step {global_step}/{total_steps} " f"(epoch {epoch+1}) " f"loss={loss.item():.4f} avg={avg:.4f} " f"lr={lr_now:.2e} vram={vram:.1f}GB " f"{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}, ckpt_path) except RuntimeError as e: if "out of memory" in str(e): torch.cuda.empty_cache() optimizer.zero_grad() global_step += 1 scheduler.step() continue raise # Save final final_path = os.path.join(OUTPUT_DIR, "head_final.pth") torch.save(head.state_dict(), final_path) print(f"\nSaved: {final_path}") total_time = time.time() - t0 print(f"Training complete: {total_steps} steps, {total_time/3600:.1f} hours") log_file.close() if __name__ == "__main__": train()