Stage 4C: direct classifier-score supervision, F1 0.729 (+0.006 over 4B)

stage_4c/README.md

@@ -0,0 +1,44 @@
+# Stage 4C: Direct Classifier-Score Supervision
+
+Same 3.27 M student as Stage 4. Same 40-D output. Different loss:
+
+```python
+student_score = student_out[pos_dims].sum() - student_out[neg_dims].sum()
+teacher_score = teacher_target[pos_dims].sum() - teacher_target[neg_dims].sum()
+loss = (student_score - teacher_score) ** 2
+```
+
+The student is optimized to match the teacher's scalar classifier output, not the 768-D feature vector (Stage 4B) or the 40 individual dims (Stage 4A).
+
+## Result
+
+```
+Stage   Student params   Loss                              F1      Threshold
+ 4       3.27 M          MSE on 40-D per-dim              0.710    26.3
+ 4B     15.67 M          cosine on 768-D                  0.723   168.0    (scale drifted)
+ 4C      3.27 M          MSE on scalar sum-difference     0.729    25.8    (matches teacher 25.3)
+ 0      85.64 M (ViT-B)  baseline                         0.889    25.3
+```
+
+Threshold converged to 25.84 — almost exactly the teacher's 25.28. The scale calibration works as designed.
+
+F1 improved by only +0.006 over Stage 4B. All three student experiments plateau around 0.72-0.73 with high recall (≥0.95) and precision ~0.58. The student converges on an "over-fire" operating point that no amount of loss-shape tuning fixes.
+
+## What this says
+
+The bottleneck is not loss choice or target geometry but the student's ability to learn the underlying scene-level signal at this scale. Closing the F1 gap to baseline 0.889 at the 3 M parameter tier probably requires:
+
+- Stronger image augmentation (mosaic, color jitter, rand-augment)
+- Warm-starting from a pre-trained backbone (EUPE-ViT-T already distilled) rather than from scratch
+- More training data beyond COCO-only (117 K images is tight for a specialist from scratch)
+
+Parameter scaling alone doesn't help; loss reshape alone doesn't help. Data and initialization are the remaining knobs.
+
+## Files
+
+- `train.py` — training loop (direct scalar MSE)
+- `student_ep{5,10,15}.safetensors` — intermediate checkpoints
+- `student_final.safetensors` — final weights
+- `training_log.json` — per-epoch loss + F1
+
+Uses the same `student.py` as Stage 4.

stage_4c/student_ep10.safetensors

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

stage_4c/student_ep15.safetensors

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

stage_4c/student_ep5.safetensors

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

stage_4c/student_final.safetensors

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

stage_4c/train.py

@@ -0,0 +1,189 @@
+"""Stage 4C: direct classifier-score supervision.
+
+Same 3.27M student architecture as Stage 4. Same 40-D output. But the loss
+is on the *classifier score* rather than the per-dim values:
+
+    student_score = student_out[pos_dims].sum() - student_out[neg_dims].sum()
+    teacher_score = teacher_target[pos_dims].sum() - teacher_target[neg_dims].sum()
+    loss = (student_score - teacher_score) ** 2
+
+The student is optimized to produce the same binary decision as the teacher
+at the classifier threshold, not to reproduce the teacher's feature geometry
+dim-by-dim. If the Stage 4B plateau at F1 0.723 was caused by even
+small per-dim errors accumulating into scalar miscalibration, this should
+close the gap.
+"""
+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, '/mnt/d/_tmp/1pc_repo/stage_4')
+from student import SpecialistStudent
+
+COCO_ROOT = '/home/zootest/datasets/coco'
+TARGETS = f'{COCO_ROOT}/stage4_teacher_targets/targets.pt'
+CLASSIFIER = '/mnt/d/_tmp/1pc_repo/stage_0/classifier.json'
+OUT_DIR = '/mnt/d/_tmp/1pc_repo/stage_4c'
+DEVICE = 'cuda'
+RES = 768
+BATCH = 16
+LR = 3e-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']         # (N, 40)
+        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()         # (40,)
+        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 eval_f1(student, pos_idx, neg_idx):
+    coco = COCO(f'{COCO_ROOT}/annotations/instances_val2017.json')
+    img_ids = sorted(coco.getImgIds())[:500]
+    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()[0]
+            scores.append((out[pos_idx].sum() - out[neg_idx].sum()).item())
+            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)
+    pack = torch.load(TARGETS, map_location='cpu', weights_only=False)
+    print(f'[init] {pack["targets"].shape[0]} targets  shape {tuple(pack["targets"].shape)}',
+          flush=True)
+
+    # In the 40-D target vector, [0..19] are pos dims, [20..39] are neg dims (built that way by prepare_targets)
+    pos_idx = torch.arange(0, 20, device=DEVICE)
+    neg_idx = torch.arange(20, 40, device=DEVICE)
+
+    # Pre-compute teacher scalar scores: (N,)
+    teacher_scalar = pack['targets'].float()[:, :20].sum(1) - pack['targets'].float()[:, 20:].sum(1)
+    pack['teacher_scalar'] = teacher_scalar
+    print(f'[init] teacher scalar stats: mean={teacher_scalar.mean():.3f}  '
+          f'std={teacher_scalar.std():.3f}', 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 = SpecialistStudent().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': 'MSE_on_classifier_scalar', '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)                              # (B, 40)
+            pred = pred.float()
+            student_scalar = pred[:, :20].sum(1) - pred[:, 20:].sum(1)    # (B,)
+            teacher_scalar_b = y[:, :20].sum(1) - y[:, 20:].sum(1)
+            loss = F.mse_loss(student_scalar, teacher_scalar_b)
+            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(student, pos_idx, neg_idx)
+        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_4c/training_log.json

@@ -0,0 +1,126 @@
+{
+  "student_params": 3267304,
+  "loss": "MSE_on_classifier_scalar",
+  "epochs": [
+    {
+      "epoch": 1,
+      "loss": 338.8909246066799,
+      "F1": 0.7214378118515015,
+      "precision": 0.5653923749923706,
+      "recall": 0.9964538812637329,
+      "threshold": 24.078125
+    },
+    {
+      "epoch": 2,
+      "loss": 328.9103960763321,
+      "F1": 0.7205128073692322,
+      "precision": 0.564257025718689,
+      "recall": 0.9964538812637329,
+      "threshold": 28.883697509765625
+    },
+    {
+      "epoch": 3,
+      "loss": 328.8529856922911,
+      "F1": 0.72258061170578,
+      "precision": 0.5679513216018677,
+      "recall": 0.9929078221321106,
+      "threshold": 28.009765625
+    },
+    {
+      "epoch": 4,
+      "loss": 326.7125272106021,
+      "F1": 0.7221510410308838,
+      "precision": 0.5651302337646484,
+      "recall": 1.0,
+      "threshold": 26.10400390625
+    },
+    {
+      "epoch": 5,
+      "loss": 325.2882255967271,
+      "F1": 0.7241829633712769,
+      "precision": 0.5734989643096924,
+      "recall": 0.9822695255279541,
+      "threshold": 26.7830810546875
+    },
+    {
+      "epoch": 6,
+      "loss": 326.14567763865386,
+      "F1": 0.7298701405525208,
+      "precision": 0.5758196711540222,
+      "recall": 0.9964538812637329,
+      "threshold": 26.697509765625
+    },
+    {
+      "epoch": 7,
+      "loss": 325.15346100816646,
+      "F1": 0.7221510410308838,
+      "precision": 0.5651302337646484,
+      "recall": 1.0,
+      "threshold": 24.860595703125
+    },
+    {
+      "epoch": 8,
+      "loss": 321.891113616252,
+      "F1": 0.7402032017707825,
+      "precision": 0.6265356540679932,
+      "recall": 0.9042553305625916,
+      "threshold": 25.110595703125
+    },
+    {
+      "epoch": 9,
+      "loss": 324.80503095718103,
+      "F1": 0.7258687615394592,
+      "precision": 0.5696969628334045,
+      "recall": 1.0,
+      "threshold": 24.62255859375
+    },
+    {
+      "epoch": 10,
+      "loss": 323.8681324547888,
+      "F1": 0.7338129878044128,
+      "precision": 0.6174334287643433,
+      "recall": 0.9042553305625916,
+      "threshold": 25.044189453125
+    },
+    {
+      "epoch": 11,
+      "loss": 322.2317366366947,
+      "F1": 0.7270233035087585,
+      "precision": 0.5928411483764648,
+      "recall": 0.9397163391113281,
+      "threshold": 25.876220703125
+    },
+    {
+      "epoch": 12,
+      "loss": 323.8775074345158,
+      "F1": 0.7272727489471436,
+      "precision": 0.5786163806915283,
+      "recall": 0.978723406791687,
+      "threshold": 26.55078125
+    },
+    {
+      "epoch": 13,
+      "loss": 322.46046450084134,
+      "F1": 0.729658842086792,
+      "precision": 0.5791666507720947,
+      "recall": 0.9858155846595764,
+      "threshold": 25.8525390625
+    },
+    {
+      "epoch": 14,
+      "loss": 321.4742774784218,
+      "F1": 0.7279894948005676,
+      "precision": 0.5782880783081055,
+      "recall": 0.9822695255279541,
+      "threshold": 25.83251953125
+    },
+    {
+      "epoch": 15,
+      "loss": 322.15041381942024,
+      "F1": 0.7289474010467529,
+      "precision": 0.5794979333877563,
+      "recall": 0.9822695255279541,
+      "threshold": 25.84033203125
+    }
+  ]
+}