stage_4b/train.py

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