Stage 1: output-channel pruning (Path A, exact parity with Stage 0)

stage_1/README.md

@@ -1,5 +1,32 @@
 # Stage 1: Output-Channel Pruning
 
-Reserved. See repo root README for plan.
+Path A implementation: runtime slicing with classifier fusion. Zero behavior drift from Stage 0.
 
-Scope: keep only the 100 feature dimensions the Stage 0 classifier reads, remove the remaining 668 output channels from EUPE-ViT-B's final projection. No retraining. Expected to preserve F1 exactly.
+## What changed
+
+The backbone emits a 768-D vector per token. Stage 1 wraps that output so downstream code sees only the 100 dimensions the classifier reads. The classifier is fused into a single `Linear(100, 1)` layer with ternary {+1, -1} fixed weights and one free parameter expressed as the negated threshold bias.
+
+```python
+from model import Stage1PersonClassifier
+model = Stage1PersonClassifier.from_pretrained_argus('phanerozoic/argus')
+score, pred = model(image_tensor)   # (B,) float, (B,) bool
+```
+
+## What did not change
+
+Compute: identical to Stage 0. The backbone still produces all 768 output channels, the final LayerNorm still runs across 768 dimensions, and the 668 unused channels are sliced off at the end. This stage exists to crystallize the interface (single classifier head, single learnable scalar) before later stages actually shrink the backbone.
+
+## Evaluation
+
+5000 COCO val 2017 images, live Argus forward pass at 768 pixel input:
+
+```
+Stage 0 F1      0.8886
+Stage 1 F1      0.8886    (exact parity, match=true)
+```
+
+See `eval.json` for precision and recall.
+
+## What Path B would look like
+
+Path B drops the last block's MLP `fc2` output rows for the 668 unused dimensions, calibrates the final LayerNorm using fixed per-channel statistics collected from a corpus, and reduces `fc2` from (3072 → 768) to (3072 → 100). Saves ~2.3M backbone parameters (1.6% of 85M). Expected F1 drift is small (<0.02) but non-zero due to the LayerNorm approximation. Not implemented in this stage; belongs in a later iteration or a Stage 1B branch if pursued.

stage_1/eval.json

@@ -0,0 +1,9 @@
+{
+  "stage": 1,
+  "F1": 0.8885542750358582,
+  "precision": 0.9011073112487793,
+  "recall": 0.8763461112976074,
+  "threshold": 25.284494400024414,
+  "stage_0_baseline_F1": 0.8886,
+  "match": true
+}

stage_1/model.py

@@ -0,0 +1,88 @@
+"""Stage 1: output-channel pruning.
+
+The full 768-D EUPE-ViT-B output token is produced as before, the final
+LayerNorm runs as before, and then only the 100 dimensions the classifier
+reads are retained. The classifier is fused into a single Linear(100, 1)
+layer with ternary {+1, 0, -1} fixed weights and one free bias (threshold
+expressed as negative bias). Inference is identical to Stage 0 by
+construction. No compute savings; this stage just cleans the interface
+and sets up the weight shapes that later stages will attack.
+
+Usage:
+    model = Stage1PersonClassifier.from_pretrained_argus('phanerozoic/argus')
+    score, pred = model(image_tensor)
+"""
+import json, os, sys
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Stage1PersonClassifier(nn.Module):
+    """EUPE-ViT-B -> 100 dim slice -> ternary linear head -> binary decision.
+
+    pos_dims and neg_dims are index tensors into the 768-D output. The
+    classifier weight matrix stored in `retained_weight` has shape (1, 100)
+    where positive-dim positions are +1 and negative-dim positions are -1.
+    Bias equals the negated threshold.
+    """
+
+    def __init__(self, argus_model, pos_dims, neg_dims, threshold):
+        super().__init__()
+        self.backbone = argus_model.backbone
+        retained = list(pos_dims) + list(neg_dims)
+        self.register_buffer('retained_dims', torch.tensor(retained, dtype=torch.long))
+        w = torch.zeros(1, len(retained))
+        w[0, : len(pos_dims)] = 1.0
+        w[0, len(pos_dims):] = -1.0
+        self.register_buffer('retained_weight', w)
+        # Stored as a free parameter so gradient descent could retune.
+        self.threshold = nn.Parameter(torch.tensor(float(threshold)))
+        self.D = 768
+
+    @torch.inference_mode()
+    def forward(self, x):
+        """x: (B, 3, 768, 768) normalized (ImageNet stats).
+
+        Returns (score, pred) where score is (B,) float and pred is (B,) bool.
+        """
+        with torch.autocast('cuda', dtype=torch.bfloat16):
+            out = self.backbone.forward_features(x)
+        patches = out['x_norm_patchtokens'].float()   # (B, 2304, 768)
+        ln = F.layer_norm(patches, [self.D])
+        pooled = ln.max(dim=1).values                  # (B, 768)
+        retained = pooled.index_select(-1, self.retained_dims)  # (B, 100)
+        score = F.linear(retained, self.retained_weight).squeeze(-1)  # (B,)
+        pred = score > self.threshold
+        return score, pred
+
+    @classmethod
+    def from_pretrained_argus(cls, repo_or_path='phanerozoic/argus',
+                              classifier_json='classifier.json'):
+        """Load Argus, read classifier.json, build the wrapper."""
+        from transformers import AutoModel
+        argus = AutoModel.from_pretrained(repo_or_path, trust_remote_code=True)
+        with open(classifier_json) as f:
+            c = json.load(f)
+        return cls(argus, c['pos_dims'], c['neg_dims'], c['threshold'])
+
+
+if __name__ == '__main__':
+    # Smoke test
+    from transformers import AutoModel
+    argus = AutoModel.from_pretrained('/mnt/d/Argus', trust_remote_code=True)
+    c = json.load(open(Path(__file__).parent / '..' / 'stage_0' / 'classifier.json'))
+    m = Stage1PersonClassifier(argus, c['pos_dims'], c['neg_dims'], c['threshold'])
+    m = m.cuda().eval()
+    n_all = sum(p.numel() for p in m.parameters())
+    n_backbone = sum(p.numel() for p in m.backbone.parameters())
+    n_head = n_all - n_backbone
+    print(f'total params: {n_all:,}')
+    print(f'backbone params: {n_backbone:,}')
+    print(f'head params: {n_head} (one learnable threshold; weights are fixed buffers)')
+
+    x = torch.randn(2, 3, 768, 768, device='cuda')
+    score, pred = m(x)
+    print(f'forward OK. score={score.tolist()}  pred={pred.tolist()}')

stage_1/verify.py

@@ -0,0 +1,72 @@
+"""Verify Stage 1 matches Stage 0 F1 exactly on COCO val 2017."""
+import os, sys, json, time
+import torch
+import numpy as np
+from PIL import Image
+from pycocotools.coco import COCO
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from model import Stage1PersonClassifier
+from transformers import AutoModel
+
+COCO_ROOT = '/home/zootest/datasets/coco'
+CLASSIFIER_JSON = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                               '..', 'stage_0', 'classifier.json')
+RES = 768
+
+
+def main():
+    print('[init] loading Argus + classifier', flush=True)
+    argus = AutoModel.from_pretrained('/mnt/d/Argus', trust_remote_code=True)
+    c = json.load(open(CLASSIFIER_JSON))
+    m = Stage1PersonClassifier(argus, c['pos_dims'], c['neg_dims'], c['threshold']).cuda().eval()
+
+    coco = COCO(f'{COCO_ROOT}/annotations/instances_val2017.json')
+    img_ids = sorted(coco.getImgIds())
+    labels = [any(a['category_id'] == 1 for a in coco.loadAnns(coco.getAnnIds(imgIds=i, iscrowd=False)))
+              for i in img_ids]
+
+    MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).cuda()
+    STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).cuda()
+
+    scores = torch.zeros(len(img_ids), device='cuda')
+    t0 = time.time()
+    for i, img_id in enumerate(img_ids):
+        info = coco.loadImgs(img_id)[0]
+        path = f"{COCO_ROOT}/val2017/{info['file_name']}"
+        img = Image.open(path).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).cuda().float() / 255.0
+        x = (x - MEAN) / STD
+        s, _ = m(x)
+        scores[i] = s[0]
+        if (i + 1) % 500 == 0:
+            print(f'  {i+1}/{len(img_ids)}  {(i+1)/(time.time()-t0):.1f} img/s', flush=True)
+
+    y = torch.tensor(labels, dtype=torch.bool, device='cuda')
+    pred = scores > m.threshold
+    tp = (pred & y).sum().float()
+    fp = (pred & ~y).sum().float()
+    fn = (~pred & y).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)
+    print(f'\n[verify] F1={f1:.4f}  P={prec:.4f}  R={rec:.4f}', flush=True)
+    print(f'[parity] Stage 0 baseline = 0.8886; Stage 1 should be identical', flush=True)
+
+    out = {
+        'stage': 1,
+        'F1': float(f1),
+        'precision': float(prec),
+        'recall': float(rec),
+        'threshold': float(m.threshold.item()),
+        'stage_0_baseline_F1': 0.8886,
+        'match': bool(abs(float(f1) - 0.8886) < 1e-3),
+    }
+    with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'eval.json'), 'w') as f:
+        json.dump(out, f, indent=2)
+    print(f'[done] wrote eval.json  match={out["match"]}', flush=True)
+
+
+if __name__ == '__main__':
+    main()