code/src/enhancements/arcface_head.py

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