stage_4c/train.py

"""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()