Stage 4B: 15.67M student + cosine loss on 768-D, F1 0.723 (+0.013 over Stage 4)

stage_4b/README.md

@@ -0,0 +1,42 @@
+# Stage 4B: Larger Specialist with Cosine Loss
+
+Tried the natural next knobs on Stage 4's specialist student: 5× bigger model, cosine similarity loss on the full 768-D pooled teacher output, longer schedule.
+
+## Setup
+
+- **Architecture**: depth 8, embed 384, 6 heads, MLP ratio 4, patch 16 → **15.67M parameters**
+- **Target**: full 768-D pooled layernormed output from EUPE-ViT-B (not the 40-dim subset used in Stage 4)
+- **Loss**: 1 − cosine_similarity(student_output, teacher_target)
+- **Schedule**: 15 epochs × 117,266 COCO train images, batch 16, AdamW lr 5e-4, cosine schedule with 3 % warmup
+- **Eval**: apply Stage 0 classifier weights to the 40 classifier-relevant dims of the student's 768-D output; sweep threshold
+
+## Result
+
+```
+Stage   Student params   Loss                F1
+ 4        3.27 M          MSE on 40-D         0.710
+ 4B      15.67 M          cosine on 768-D     0.723
+ 0       85.64 M (ViT-B)  baseline            0.889
+```
+
+Cosine loss converged in epoch 1 (0.072 → 0.061) and stayed flat through epoch 15. F1 plateau'd at 0.723 on the 500-image COCO val subset.
+
+## What this says
+
+The student reproduces the teacher's pooled feature geometry well in aggregate (cosine ≈ 0.94 across 768 dims), but the 40 classifier-relevant dims are not all equal. Even a small average error on those specific axes destroys Stage 0's precision — every epoch shows precision around 0.57 and recall approaching 1.0, i.e., the student is consistently over-firing.
+
+Two candidate next iterations:
+
+1. **Dim-weighted cosine**: scale the cosine loss by a per-dim importance weight, with the 40 classifier-relevant dims weighted heavily. The student would then be forced to reproduce those exact values rather than any 40 dims of equal average fidelity.
+2. **Direct classifier supervision**: train the student to minimize `|score_student - score_teacher|` where `score = sum(pos_dims) - sum(neg_dims)`, not the 768-D vector.
+
+Either is cheaper than further capacity/epoch scaling.
+
+## Files
+
+- `student.py` — architecture
+- `prepare_targets_768.py` — builds the 768-D teacher target tensor from the ViT-B cache
+- `train.py` — training loop
+- `student_ep{5,10,15}.safetensors` — intermediate checkpoints
+- `student_final.safetensors` — final weights
+- `training_log.json` — per-epoch loss + F1

stage_4b/prepare_targets_768.py

@@ -0,0 +1,44 @@
+"""Stage 4B: build the 768-D teacher-target tensor from the ViT-B COCO cache.
+
+Per image: layernorm 768 channels across patches, max-pool patches → one
+768-D vector. Save as (N, 768) float16 + aligned img_ids.
+
+Output: /home/zootest/datasets/coco/stage4b_teacher_targets/targets.pt
+"""
+import os, glob, time
+import torch
+import torch.nn.functional as F
+
+COCO_ROOT = '/home/zootest/datasets/coco'
+CACHE = f'{COCO_ROOT}/feature_cache_768'
+OUT_DIR = f'{COCO_ROOT}/stage4b_teacher_targets'
+D = 768
+
+
+def main():
+    os.makedirs(OUT_DIR, exist_ok=True)
+    shards = sorted(glob.glob(f'{CACHE}/shard_*.pt'))
+    print(f'[init] {len(shards)} shards', flush=True)
+
+    targets, img_ids = [], []
+    t0 = time.time()
+    for si, sp in enumerate(shards):
+        shard = torch.load(sp, map_location='cpu', weights_only=False)
+        for entry in shard:
+            feat = entry['spatial'].float()                                # (768, 48, 48)
+            ln = F.layer_norm(feat.permute(1, 2, 0).reshape(-1, D), [D])   # (2304, 768)
+            pooled = ln.max(dim=0).values                                   # (768,)
+            targets.append(pooled.half())
+            img_ids.append(int(entry.get('img_id', len(img_ids))))
+        rate = (si + 1) / max(time.time() - t0, 1e-6)
+        print(f'  shard {si+1}/{len(shards)}  n={len(targets)}  '
+              f'{(time.time()-t0):.0f}s  ETA {(len(shards)-si-1)/rate:.0f}s', flush=True)
+
+    targets = torch.stack(targets)
+    torch.save({'targets': targets, 'img_ids': img_ids},
+               f'{OUT_DIR}/targets.pt')
+    print(f'[done] {targets.shape} -> {OUT_DIR}/targets.pt', flush=True)
+
+
+if __name__ == '__main__':
+    main()

stage_4b/student.py

@@ -0,0 +1,68 @@
+"""Stage 4B bigger specialist student.
+
+Depth 8, embed 384, 6 heads, MLP ratio 4. Emits a 768-D vector per image
+(matches the full EUPE-ViT-B pooled layernormed output) for cosine-similarity
+distillation.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, in_ch=3, embed_dim=384, patch_size=16):
+        super().__init__()
+        self.proj = nn.Conv2d(in_ch, embed_dim, patch_size, stride=patch_size)
+
+    def forward(self, x):
+        x = self.proj(x)
+        return x.flatten(2).transpose(1, 2)
+
+
+class Block(nn.Module):
+    def __init__(self, dim, heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
+        self.norm2 = nn.LayerNorm(dim)
+        hidden = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(), nn.Linear(hidden, dim))
+
+    def forward(self, x):
+        h = self.norm1(x)
+        h, _ = self.attn(h, h, h, need_weights=False)
+        x = x + h
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class Stage4BStudent(nn.Module):
+    """Outputs a 768-D vector per image to match the EUPE-ViT-B teacher."""
+    def __init__(self, out_dim=768, embed_dim=384, depth=8, heads=6,
+                 patch_size=16, img_size=768, mlp_ratio=4.0):
+        super().__init__()
+        self.patch = PatchEmbed(3, embed_dim, patch_size)
+        self.num_patches = (img_size // patch_size) ** 2
+        self.pos = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim))
+        nn.init.trunc_normal_(self.pos, std=0.02)
+        self.blocks = nn.ModuleList([Block(embed_dim, heads, mlp_ratio) for _ in range(depth)])
+        self.norm = nn.LayerNorm(embed_dim)
+        self.head = nn.Linear(embed_dim, out_dim)
+
+    def forward(self, x):
+        tokens = self.patch(x)
+        tokens = tokens + self.pos[:, :tokens.shape[1]]
+        for blk in self.blocks:
+            tokens = blk(tokens)
+        tokens = self.norm(tokens)
+        pooled = tokens.max(dim=1).values
+        return self.head(pooled)
+
+
+if __name__ == '__main__':
+    m = Stage4BStudent()
+    total = sum(p.numel() for p in m.parameters())
+    print(f'Stage 4B student: {total:,} params = {total/1e6:.2f}M')
+    x = torch.randn(2, 3, 768, 768)
+    y = m(x)
+    print(f'forward OK  output shape: {tuple(y.shape)}')

stage_4b/student_ep10.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e29f518c514d0ea26b5340104e054273e48a2a74ee7db025ce47c42cc28b523
+size 62698160

stage_4b/student_ep15.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd541b1fce0fe10af77471404fa8743960f10fb1798e4ebc0409437f30cdb83f
+size 62698160

stage_4b/student_ep5.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4be68d45bc8922a2d5a34ff040b617030abbb505071463b684c83c4f67f7b1bb
+size 62698160

stage_4b/student_final.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd541b1fce0fe10af77471404fa8743960f10fb1798e4ebc0409437f30cdb83f
+size 62698160

stage_4b/train.py

@@ -0,0 +1,176 @@
+"""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()

stage_4b/training_log.json

@@ -0,0 +1,127 @@
+{
+  "student_params": 15672192,
+  "loss": "cosine_1_minus_sim",
+  "target_dim": 768,
+  "epochs": [
+    {
+      "epoch": 1,
+      "loss": 0.07253991591294745,
+      "F1": 0.7222222089767456,
+      "precision": 0.5759493708610535,
+      "recall": 0.9680851101875305,
+      "threshold": 77.3359375
+    },
+    {
+      "epoch": 2,
+      "loss": 0.06178997803267662,
+      "F1": 0.7249357104301453,
+      "precision": 0.5685483813285828,
+      "recall": 1.0,
+      "threshold": 107.1875
+    },
+    {
+      "epoch": 3,
+      "loss": 0.0617156123302943,
+      "F1": 0.7195902466773987,
+      "precision": 0.563126266002655,
+      "recall": 0.9964538812637329,
+      "threshold": 148.1875
+    },
+    {
+      "epoch": 4,
+      "loss": 0.061347012896584736,
+      "F1": 0.7253885865211487,
+      "precision": 0.5714285969734192,
+      "recall": 0.9929078221321106,
+      "threshold": 153.75
+    },
+    {
+      "epoch": 5,
+      "loss": 0.06122595644468568,
+      "F1": 0.7237353920936584,
+      "precision": 0.5705521702766418,
+      "recall": 0.9893617033958435,
+      "threshold": 159.734375
+    },
+    {
+      "epoch": 6,
+      "loss": 0.06132089458562277,
+      "F1": 0.7195902466773987,
+      "precision": 0.563126266002655,
+      "recall": 0.9964538812637329,
+      "threshold": 156.5625
+    },
+    {
+      "epoch": 7,
+      "loss": 0.061309009904699646,
+      "F1": 0.7220779061317444,
+      "precision": 0.5696721076965332,
+      "recall": 0.9858155846595764,
+      "threshold": 174.6875
+    },
+    {
+      "epoch": 8,
+      "loss": 0.06113341519135151,
+      "F1": 0.720720648765564,
+      "precision": 0.5656565427780151,
+      "recall": 0.9929078221321106,
+      "threshold": 168.375
+    },
+    {
+      "epoch": 9,
+      "loss": 0.060983790468551714,
+      "F1": 0.7205128073692322,
+      "precision": 0.564257025718689,
+      "recall": 0.9964538812637329,
+      "threshold": 168.28125
+    },
+    {
+      "epoch": 10,
+      "loss": 0.061084045586111,
+      "F1": 0.7260638475418091,
+      "precision": 0.5808510780334473,
+      "recall": 0.9680851101875305,
+      "threshold": 165.9375
+    },
+    {
+      "epoch": 11,
+      "loss": 0.061100886026898504,
+      "F1": 0.7221510410308838,
+      "precision": 0.5651302337646484,
+      "recall": 1.0,
+      "threshold": 171.53125
+    },
+    {
+      "epoch": 12,
+      "loss": 0.06096925960034696,
+      "F1": 0.7195902466773987,
+      "precision": 0.563126266002655,
+      "recall": 0.9964538812637329,
+      "threshold": 168.1875
+    },
+    {
+      "epoch": 13,
+      "loss": 0.060978461868108035,
+      "F1": 0.7253613471984863,
+      "precision": 0.5762004256248474,
+      "recall": 0.978723406791687,
+      "threshold": 167.40625
+    },
+    {
+      "epoch": 14,
+      "loss": 0.060976944405298966,
+      "F1": 0.7216494679450989,
+      "precision": 0.5668016076087952,
+      "recall": 0.9929078221321106,
+      "threshold": 168.125
+    },
+    {
+      "epoch": 15,
+      "loss": 0.06102545032333958,
+      "F1": 0.7230768799781799,
+      "precision": 0.5662650465965271,
+      "recall": 1.0,
+      "threshold": 168.09375
+    }
+  ]
+}