#!/usr/bin/env python3 """ArcFace margin-based head (Deng CVPR 2019). Frozen backbone + ArcFace loss: cos(θ+m) instead of cos(θ) for true class. Enforces angular margin → better separation, especially for tail classes (each class gets its own angular region). """ import os, sys, json, argparse, time, math 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 ArcFaceHead(nn.Module): def __init__(self, feat_dim, n_classes, s=30.0, m=0.30): super().__init__() self.W = nn.Parameter(torch.randn(n_classes, feat_dim)) nn.init.xavier_uniform_(self.W) self.s = s; self.m = m self.n_classes = n_classes self.cos_m = math.cos(m) self.sin_m = math.sin(m) self.th = math.cos(math.pi - m) self.mm = math.sin(math.pi - m) * m def forward(self, feats, target=None): f = F.normalize(feats.float(), dim=-1) w = F.normalize(self.W.float(), dim=-1) cos = (f @ w.t()).clamp(-1 + 1e-7, 1 - 1e-7) if target is None: return cos * self.s sin = torch.sqrt((1.0 - cos * cos).clamp_min(1e-7)) cos_tm = cos * self.cos_m - sin * self.sin_m cos_tm = torch.where(cos > self.th, cos_tm, cos - self.mm) one_hot = F.one_hot(target, self.n_classes).float() logits = (one_hot * cos_tm + (1.0 - one_hot) * cos) * self.s return logits @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=3e-3) ap.add_argument('--s', type=float, default=30.0) ap.add_argument('--margin', type=float, default=0.2) ap.add_argument('--warmup-epochs', type=int, default=2) ap.add_argument('--grad-clip', type=float, default=5.0) 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 = ArcFaceHead(feat_dim, n_cls, s=args.s, m=args.margin).to(device) opt = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=1e-4) warm = torch.optim.lr_scheduler.LinearLR(opt, start_factor=0.1, end_factor=1.0, total_iters=max(1, args.warmup_epochs)) cos = torch.optim.lr_scheduler.CosineAnnealingLR( opt, T_max=max(1, args.epochs - args.warmup_epochs)) sched = torch.optim.lr_scheduler.SequentialLR( opt, schedulers=[warm, cos], milestones=[args.warmup_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) logits = head(f, target=y) loss = F.cross_entropy(logits, y, label_smoothing=0.05) if not torch.isfinite(loss): log(f"non-finite loss at ep {ep}; skip batch") opt.zero_grad(); continue opt.zero_grad(); loss.backward() torch.nn.utils.clip_grad_norm_(head.parameters(), args.grad_clip) 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) logits = head(f) # no margin at eval 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, 's': args.s, 'margin': args.margin}, args.output) log(f"=== ArcFace BEST: {best:.4f}") if __name__ == '__main__': main()