| """ |
| 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)} |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
| def train(): |
| os.makedirs(OUTPUT_DIR, exist_ok=True) |
|
|
| print("=" * 60) |
| print("Full COCO training: Cofiber Threshold (S)") |
| print("=" * 60) |
|
|
| |
| 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 |
|
|
| |
| 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)") |
|
|
| |
| 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 = 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] |
| 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")) |
|
|
| |
| 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: |
| |
| 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, 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 |
|
|
| |
| 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() |
|
|
