#!/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()