code/src/enhancements/pcl_head.py

#!/usr/bin/env python3
"""Prototypical Contrastive Learning (PCL, Li ICLR 2021).

Frozen backbone + trainable projector + learnable class prototypes.
Loss = supervised InfoNCE (positives: same class) + prototype alignment.

Simpler than full MoCo; no momentum encoder. Treats each batch mini as
"queue" — feasible given our batch sizes 128+.
"""
import os, sys, json, argparse, time
from pathlib import Path
from collections import Counter

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, WeightedRandomSampler
from PIL import Image

ROOT = Path("/arf/scratch/stakan/hitit-proje")
sys.path.insert(0, str(ROOT / "hitit_ocr/src"))
from train_classification import build_backbone, get_arch_img_size, HititClsDataset

def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)


class Projector(nn.Module):
    def __init__(self, in_dim, hid=512, out=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hid), nn.GELU(),
            nn.Linear(hid, out),
        )
    def forward(self, x): return self.net(x)


class PCLHead(nn.Module):
    def __init__(self, feat_dim, n_classes, proj_dim=128, tau=0.1):
        super().__init__()
        self.proj = Projector(feat_dim, hid=512, out=proj_dim)
        self.proto = nn.Parameter(F.normalize(torch.randn(n_classes, proj_dim), dim=-1))
        self.cls = nn.Linear(feat_dim, n_classes)  # parallel linear for CE
        self.tau = tau

    def forward(self, feats):
        z = F.normalize(self.proj(feats), dim=-1)
        proto = F.normalize(self.proto, dim=-1)
        sim = z @ proto.t() / self.tau  # (B, C)
        logits_cls = self.cls(feats)
        return sim, logits_cls, z


def supcon_loss(z, labels, tau=0.1):
    sim = z @ z.t() / tau
    mask = (labels.view(-1, 1) == labels.view(1, -1)).float()
    # remove self
    B = z.size(0)
    mask.fill_diagonal_(0)
    sim.fill_diagonal_(-1e4)
    logp = sim - sim.logsumexp(-1, keepdim=True)
    n_pos = mask.sum(-1).clamp_min(1)
    return -(mask * logp).sum(-1).div(n_pos).mean()


@torch.no_grad()
def extract(model, x, dtype):
    raw = model.module if hasattr(model, 'module') else model
    with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
        if hasattr(raw, 'forward_features'):
            f = raw.forward_features(x)
            if f.dim() == 3: f = f[:, 0]
            elif f.dim() == 4: f = f.mean(dim=(2, 3))
        else:
            f = raw(x)
    return f.float()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', required=True)
    ap.add_argument('--manifest', required=True)
    ap.add_argument('--val-fold', type=int, default=0)
    ap.add_argument('--min-samples', type=int, default=10)
    ap.add_argument('--epochs', type=int, default=30)
    ap.add_argument('--lr', type=float, default=5e-3)
    ap.add_argument('--tau', type=float, default=0.1)
    ap.add_argument('--supcon-weight', type=float, default=0.5)
    ap.add_argument('--proto-weight', type=float, default=0.3)
    ap.add_argument('--batch-size', type=int, default=128)
    ap.add_argument('--output', required=True)
    args = ap.parse_args()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32

    arch, path = args.ckpt.split(':', 1)
    ck = torch.load(path, map_location='cpu', weights_only=False)
    label_to_idx = ck['label_to_idx']
    n_cls = len(label_to_idx)
    img_size = get_arch_img_size(arch)
    backbone = build_backbone(arch, n_classes=n_cls, img_size_override=img_size).to(device)
    sd = ck['model']
    sd = {k.replace('module.', '', 1): v for k, v in sd.items()}
    sd = {k.replace('_orig_mod.', '', 1): v for k, v in sd.items()}
    sd = {k: v for k, v in sd.items() if k != 'n_averaged'}
    backbone.load_state_dict(sd, strict=False)
    for p in backbone.parameters(): p.requires_grad = False
    backbone.eval()

    cfg = {'img_size': img_size}
    tr_ds = HititClsDataset(args.manifest, cfg, is_train=True,
                             val_fold=args.val_fold, label_to_idx=label_to_idx,
                             min_samples=args.min_samples)
    va_ds = HititClsDataset(args.manifest, cfg, is_train=False,
                             val_fold=args.val_fold, label_to_idx=label_to_idx,
                             min_samples=args.min_samples)
    cc = Counter(tr_ds.label_to_idx[r['unified_label']] for r in tr_ds.records)
    sample_w = np.array([1.0 / max(1, cc[tr_ds.label_to_idx[r['unified_label']]])**0.5
                          for r in tr_ds.records], dtype=np.float32)
    tr_dl = DataLoader(tr_ds, batch_size=args.batch_size,
                        sampler=WeightedRandomSampler(sample_w, len(tr_ds), True),
                        num_workers=6, pin_memory=True, drop_last=True)
    va_dl = DataLoader(va_ds, batch_size=args.batch_size, shuffle=False,
                        num_workers=4, pin_memory=True)

    with torch.no_grad():
        feat_dim = extract(backbone, torch.zeros(1, 3, img_size, img_size, device=device), dtype).size(-1)
    head = PCLHead(feat_dim, n_cls, tau=args.tau).to(device)
    opt = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)

    best, best_probs, best_y = 0, None, None
    for ep in range(args.epochs):
        head.train()
        tl, nb = 0, 0
        for x, y in tr_dl:
            x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
            f = extract(backbone, x, dtype)
            sim, logits_cls, z = head(f)
            loss_ce = F.cross_entropy(logits_cls, y, label_smoothing=0.1)
            loss_proto = F.cross_entropy(sim, y)  # align with prototype
            loss_sc = supcon_loss(z, y, tau=args.tau) if z.size(0) > 1 else torch.zeros(1, device=device)
            loss = loss_ce + args.proto_weight * loss_proto + args.supcon_weight * loss_sc
            opt.zero_grad(); loss.backward(); opt.step()
            tl += loss.item(); nb += 1
        sched.step()
        head.eval()
        cs, tot, probs_all, y_all = 0, 0, [], []
        with torch.no_grad():
            for x, y in va_dl:
                x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
                f = extract(backbone, x, dtype)
                sim, logits_cls, _ = head(f)
                # ensemble sim + cls
                logits = 0.5 * sim + 0.5 * logits_cls
                p = F.softmax(logits.float(), dim=-1)
                probs_all.append(p.cpu()); y_all.append(y.cpu())
                cs += (logits.argmax(-1) == y).sum().item(); tot += y.size(0)
        acc = cs / max(1, tot)
        if ep % 5 == 0 or ep == args.epochs - 1:
            log(f"ep {ep}: loss={tl/max(1,nb):.4f} val_top1={acc:.4f}")
        if acc > best:
            best = acc
            best_probs = torch.cat(probs_all); best_y = torch.cat(y_all)
            torch.save({'head_state': head.state_dict(), 'label_to_idx': label_to_idx,
                         'probs': best_probs, 'targets': best_y, 'top1': best,
                         'backbone_arch': arch},
                        args.output)
    log(f"=== PCL BEST: {best:.4f}")

if __name__ == '__main__':
    main()