"""Stage 4B training: cosine loss on full 768-D pooled teacher output, 30 epochs.""" import os, sys, time, json, math import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from pycocotools.coco import COCO from safetensors.torch import save_file HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) from student import Stage4BStudent COCO_ROOT = '/home/zootest/datasets/coco' TARGETS = f'{COCO_ROOT}/stage4b_teacher_targets/targets.pt' CLASSIFIER = '/mnt/d/_tmp/1pc_repo/stage_0/classifier.json' OUT_DIR = '/mnt/d/_tmp/1pc_repo/stage_4b' DEVICE = 'cuda' RES = 768 BATCH = 16 LR = 5e-4 WD = 1e-4 EPOCHS = 15 WARMUP_FRAC = 0.03 class CocoImgDataset(torch.utils.data.Dataset): def __init__(self, coco_root, pack): self.root = f'{coco_root}/train2017' coco = COCO(f'{coco_root}/annotations/instances_train2017.json') self.img_ids = pack['img_ids'] self.targets = pack['targets'] self.id_to_file = {i['id']: i['file_name'] for i in coco.loadImgs(coco.getImgIds())} def __len__(self): return len(self.img_ids) def __getitem__(self, i): img_id = self.img_ids[i] target = self.targets[i].float() fname = self.id_to_file.get(img_id) if fname is None: return None try: img = Image.open(f'{self.root}/{fname}').convert('RGB').resize((RES, RES), Image.BILINEAR) except Exception: return None arr = np.asarray(img, dtype=np.uint8).copy() x = torch.from_numpy(arr).permute(2, 0, 1).float() / 255.0 mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) return (x - mean) / std, target def collate(batch): batch = [b for b in batch if b is not None] if not batch: return None xs, ts = zip(*batch) return torch.stack(xs), torch.stack(ts) def cosine_loss(pred, target): # 1 - cosine similarity (per-sample mean) return (1.0 - F.cosine_similarity(pred, target, dim=-1)).mean() def eval_f1_via_stage0(student, classifier_json, n=500): with open(classifier_json) as f: c = json.load(f) pos = torch.tensor(c['pos_dims'], device=DEVICE) neg = torch.tensor(c['neg_dims'], device=DEVICE) thr = c['threshold'] coco = COCO(f'{COCO_ROOT}/annotations/instances_val2017.json') img_ids = sorted(coco.getImgIds())[:n] id_to_file = {i['id']: i['file_name'] for i in coco.loadImgs(coco.getImgIds())} MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(DEVICE) STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(DEVICE) scores, labels = [], [] student.eval() with torch.inference_mode(): for img_id in img_ids: fname = id_to_file.get(img_id) if not fname: continue img = Image.open(f'{COCO_ROOT}/val2017/{fname}').convert('RGB').resize((RES, RES), Image.BILINEAR) arr = np.asarray(img, dtype=np.uint8).copy() x = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(DEVICE).float() / 255.0 x = (x - MEAN) / STD with torch.autocast('cuda', dtype=torch.bfloat16): out = student(x).float() # (1, 768) s = (out[0, pos].sum() - out[0, neg].sum()).item() scores.append(s) labels.append(any(a['category_id'] == 1 for a in coco.loadAnns(coco.getAnnIds(imgIds=img_id, iscrowd=False)))) scores = torch.tensor(scores) labels = torch.tensor(labels, dtype=torch.bool) uniq = torch.unique(scores).sort().values best = (0, 0, 0, 0) for t in uniq.tolist()[::max(1, len(uniq) // 500)]: pred = scores > t tp = (pred & labels).sum().float() fp = (pred & ~labels).sum().float() fn = (~pred & labels).sum().float() prec = tp / (tp + fp).clamp(min=1) rec = tp / (tp + fn).clamp(min=1) f1 = (2 * prec * rec / (prec + rec).clamp(min=1e-9)).item() if f1 > best[0]: best = (f1, t, prec.item(), rec.item()) return best def main(): os.makedirs(OUT_DIR, exist_ok=True) print('[init] loading targets', flush=True) pack = torch.load(TARGETS, map_location='cpu', weights_only=False) print(f' {pack["targets"].shape[0]} teacher targets, dim {pack["targets"].shape[1]}', flush=True) ds = CocoImgDataset(COCO_ROOT, pack) loader = torch.utils.data.DataLoader( ds, batch_size=BATCH, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate, drop_last=True) student = Stage4BStudent().to(DEVICE) nparams = sum(p.numel() for p in student.parameters()) print(f'[student] {nparams:,} params = {nparams/1e6:.2f}M', flush=True) total_steps = EPOCHS * len(loader) warmup = int(total_steps * WARMUP_FRAC) opt = torch.optim.AdamW(student.parameters(), lr=LR, weight_decay=WD) sched = torch.optim.lr_scheduler.LambdaLR( opt, lambda s: s / max(1, warmup) if s < warmup else 0.5 * (1 + math.cos(math.pi * (s - warmup) / max(1, total_steps - warmup)))) log = {'student_params': nparams, 'loss': 'cosine_1_minus_sim', 'target_dim': 768, 'epochs': []} step = 0; t0 = time.time() for ep in range(EPOCHS): student.train() ep_loss, n_batches = 0.0, 0 for batch in loader: if batch is None: continue x, y = batch x = x.to(DEVICE, non_blocking=True); y = y.to(DEVICE, non_blocking=True) with torch.autocast('cuda', dtype=torch.bfloat16): pred = student(x) loss = cosine_loss(pred.float(), y) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0) opt.step(); sched.step() ep_loss += loss.item(); n_batches += 1; step += 1 if step % 500 == 0: print(f' ep {ep+1}/{EPOCHS} step {step}/{total_steps} ' f'loss={loss.item():.4f} lr={opt.param_groups[0]["lr"]:.2e} ' f'{(time.time()-t0)/60:.1f} min', flush=True) avg = ep_loss / max(1, n_batches) f1, thr, p, r = eval_f1_via_stage0(student, CLASSIFIER) print(f'[ep {ep+1}] loss={avg:.4f} F1={f1:.4f} P={p:.4f} R={r:.4f} ' f'θ={thr:.3f} {(time.time()-t0)/60:.1f} min', flush=True) log['epochs'].append({'epoch': ep + 1, 'loss': avg, 'F1': f1, 'precision': p, 'recall': r, 'threshold': thr}) if (ep + 1) % 5 == 0 or ep == EPOCHS - 1: save_file(student.state_dict(), f'{OUT_DIR}/student_ep{ep+1}.safetensors') with open(f'{OUT_DIR}/training_log.json', 'w') as f: json.dump(log, f, indent=2) save_file(student.state_dict(), f'{OUT_DIR}/student_final.safetensors') print(f'[done] total {(time.time()-t0)/60:.1f} min', flush=True) if __name__ == '__main__': main()