code/src/enhancements/linear_probe.py · savastakan/hitit-cuneiform-ocr at main

hitit-cuneiform-ocr / code /src /enhancements /linear_probe.py

Initial upload: code + 5 record checkpoints + fuse

f211247 verified 6 months ago

7.59 kB

#!/usr/bin/env python3
"""Frozen backbone + linear probe (logistic regression).
Meta DINOv3 best-practice (2025): frozen CLS + linear often beats FT for
classification, especially when downstream data is limited. Fast to run
(<1 GPU-hour on H200 for our 13k × 224 images).
"""
import os, sys, json, argparse, time
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
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

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

class SimpleDS(Dataset):
    def __init__(self, manifest, label_to_idx, tf, val_fold=0, is_val=False, min_samples=10):
        self.records = []
        from collections import Counter
        cls_count = Counter()
        with open(manifest) as f:
            for line in f:
                r = json.loads(line)
                if r.get('task') != 'classification': continue
                if not r.get('unified_label'): continue
                cls_count[r['unified_label']] += 1
        with open(manifest) as f:
            for line in f:
                r = json.loads(line)
                if r.get('task') != 'classification': continue
                if not r.get('unified_label') or r['unified_label'] not in label_to_idx: continue
                if not r.get('path') or r.get('storage') != 'fs': continue
                if r.get('integrity_ok') is False: continue
                if cls_count[r['unified_label']] < min_samples: continue
                fold = r.get('tablet_view_fold', 0)
                if is_val and fold != val_fold: continue
                if not is_val and fold == val_fold: continue
                self.records.append(r)
        self.l2i = label_to_idx; self.tf = tf
    def __len__(self): return len(self.records)
    def __getitem__(self, i):
        r = self.records[i]
        img = Image.open(r['path']).convert('RGB')
        return self.tf(img), self.l2i[r['unified_label']]

@torch.no_grad()
def extract(model, loader, device, dtype):
    model.eval()
    raw = model.module if hasattr(model, 'module') else model
    feats, ys = [], []
    for x, y in loader:
        x = x.to(device, non_blocking=True)
        with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
            if hasattr(raw, 'forward_features'):
                f = raw.forward_features(x)
                if f.dim() > 2:
                    # ViT: (B, N, D) → CLS is [0]; take CLS
                    if f.dim() == 3: f = f[:, 0]
                    else: f = f.mean(dim=tuple(range(1, f.dim()-1)))
            else:
                f = raw(x)
        feats.append(f.float().cpu()); ys.append(y)
    return torch.cat(feats), torch.cat(ys)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--backbone-arch', default='dinov3_vitl14')
    ap.add_argument('--ssl-ckpt', default=None, help='Optional SSL backbone weights')
    ap.add_argument('--manifest', required=True)
    ap.add_argument('--val-fold', type=int, default=0)
    ap.add_argument('--epochs', type=int, default=50, help='Linear head epochs')
    ap.add_argument('--lr', type=float, default=1e-2)
    ap.add_argument('--weight-decay', type=float, default=1e-4)
    ap.add_argument('--batch-size', type=int, default=256)
    ap.add_argument('--min-samples', type=int, default=10)
    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

    # Build label_to_idx from manifest (respect min_samples)
    from collections import Counter
    cls_count = Counter()
    for line in open(args.manifest):
        r = json.loads(line)
        if r.get('task') != 'classification': continue
        if not r.get('unified_label'): continue
        cls_count[r['unified_label']] += 1
    labels = sorted([l for l, n in cls_count.items() if n >= args.min_samples])
    label_to_idx = {l: i for i, l in enumerate(labels)}
    n_cls = len(label_to_idx)
    log(f"Classes (min_samples={args.min_samples}): {n_cls}")

    # Build frozen backbone (ImageNet/SSL pretrained, NO classification head training)
    model = build_backbone(args.backbone_arch, ssl_ckpt=args.ssl_ckpt, n_classes=n_cls).to(device)
    # Freeze everything
    for p in model.parameters(): p.requires_grad = False

    from torchvision import transforms
    img_size = get_arch_img_size(args.backbone_arch)
    tf = transforms.Compose([
        transforms.Resize((img_size, img_size), antialias=True),
        transforms.ToTensor(),
        transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
    ])
    tr_ds = SimpleDS(args.manifest, label_to_idx, tf, val_fold=args.val_fold, is_val=False,
                      min_samples=args.min_samples)
    va_ds = SimpleDS(args.manifest, label_to_idx, tf, val_fold=args.val_fold, is_val=True,
                      min_samples=args.min_samples)
    log(f"Train={len(tr_ds)} Val={len(va_ds)}")

    tr_dl = DataLoader(tr_ds, batch_size=args.batch_size, shuffle=False, num_workers=6)
    va_dl = DataLoader(va_ds, batch_size=args.batch_size, shuffle=False, num_workers=4)

    log("Extracting train features (frozen)...")
    tr_f, tr_y = extract(model, tr_dl, device, dtype)
    log("Extracting val features...")
    va_f, va_y = extract(model, va_dl, device, dtype)
    D = tr_f.size(-1)
    log(f"Feature dim: {D}")

    # Linear probe on extracted feats (tiny: just one Linear layer)
    head = nn.Linear(D, n_cls).to(device)
    opt = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)

    # Batch over features in memory
    tr_f_d, tr_y_d = tr_f.to(device), tr_y.to(device)
    va_f_d, va_y_d = va_f.to(device), va_y.to(device)
    best_acc = 0.0
    for ep in range(args.epochs):
        head.train()
        idx = torch.randperm(len(tr_f_d), device=device)
        for i in range(0, len(tr_f_d), 1024):
            b = idx[i:i+1024]
            logits = head(tr_f_d[b])
            loss = F.cross_entropy(logits, tr_y_d[b], label_smoothing=0.1)
            opt.zero_grad(); loss.backward(); opt.step()
        sched.step()
        # Validate
        head.eval()
        with torch.no_grad():
            logits = head(va_f_d)
            acc = (logits.argmax(-1) == va_y_d).float().mean().item()
            _, top5 = logits.topk(5, dim=-1)
            top5_acc = sum(va_y_d[i].item() in top5[i].tolist()
                           for i in range(len(va_y_d))) / len(va_y_d)
        if ep % 5 == 0 or ep == args.epochs - 1:
            log(f"ep {ep}: val_top1={acc:.4f} val_top5={top5_acc:.4f} loss={loss.item():.3f}")
        if acc > best_acc:
            best_acc = acc
            best_probs = F.softmax(head(va_f_d).float(), dim=-1).cpu()

    log(f"=== Frozen {args.backbone_arch} linear probe ===")
    log(f"  Best top-1: {best_acc:.4f}")

    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    torch.save({
        'probs': best_probs, 'targets': va_y,
        'top1': best_acc, 'label_to_idx': label_to_idx,
        'backbone': args.backbone_arch, 'feat_dim': D,
        'head_state': head.state_dict(),
    }, args.output)
    log(f"Saved → {args.output}")

if __name__ == '__main__':
    main()

	#!/usr/bin/env python3
	"""Frozen backbone + linear probe (logistic regression).

	Meta DINOv3 best-practice (2025): frozen CLS + linear often beats FT for
	classification, especially when downstream data is limited. Fast to run
	(<1 GPU-hour on H200 for our 13k × 224 images).
	"""
	import os, sys, json, argparse, time
	from pathlib import Path

	import numpy as np
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.utils.data import DataLoader, Dataset
	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

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

	class SimpleDS(Dataset):
	def __init__(self, manifest, label_to_idx, tf, val_fold=0, is_val=False, min_samples=10):
	self.records = []
	from collections import Counter
	cls_count = Counter()
	with open(manifest) as f:
	for line in f:
	r = json.loads(line)
	if r.get('task') != 'classification': continue
	if not r.get('unified_label'): continue
	cls_count[r['unified_label']] += 1
	with open(manifest) as f:
	for line in f:
	r = json.loads(line)
	if r.get('task') != 'classification': continue
	if not r.get('unified_label') or r['unified_label'] not in label_to_idx: continue
	if not r.get('path') or r.get('storage') != 'fs': continue
	if r.get('integrity_ok') is False: continue
	if cls_count[r['unified_label']] < min_samples: continue
	fold = r.get('tablet_view_fold', 0)
	if is_val and fold != val_fold: continue
	if not is_val and fold == val_fold: continue
	self.records.append(r)
	self.l2i = label_to_idx; self.tf = tf
	def __len__(self): return len(self.records)
	def __getitem__(self, i):
	r = self.records[i]
	img = Image.open(r['path']).convert('RGB')
	return self.tf(img), self.l2i[r['unified_label']]

	@torch.no_grad()
	def extract(model, loader, device, dtype):
	model.eval()
	raw = model.module if hasattr(model, 'module') else model
	feats, ys = [], []
	for x, y in loader:
	x = x.to(device, non_blocking=True)
	with torch.amp.autocast('cuda', dtype=dtype, enabled=True):
	if hasattr(raw, 'forward_features'):
	f = raw.forward_features(x)
	if f.dim() > 2:
	# ViT: (B, N, D) → CLS is [0]; take CLS
	if f.dim() == 3: f = f[:, 0]
	else: f = f.mean(dim=tuple(range(1, f.dim()-1)))
	else:
	f = raw(x)
	feats.append(f.float().cpu()); ys.append(y)
	return torch.cat(feats), torch.cat(ys)

	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument('--backbone-arch', default='dinov3_vitl14')
	ap.add_argument('--ssl-ckpt', default=None, help='Optional SSL backbone weights')
	ap.add_argument('--manifest', required=True)
	ap.add_argument('--val-fold', type=int, default=0)
	ap.add_argument('--epochs', type=int, default=50, help='Linear head epochs')
	ap.add_argument('--lr', type=float, default=1e-2)
	ap.add_argument('--weight-decay', type=float, default=1e-4)
	ap.add_argument('--batch-size', type=int, default=256)
	ap.add_argument('--min-samples', type=int, default=10)
	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

	# Build label_to_idx from manifest (respect min_samples)
	from collections import Counter
	cls_count = Counter()
	for line in open(args.manifest):
	r = json.loads(line)
	if r.get('task') != 'classification': continue
	if not r.get('unified_label'): continue
	cls_count[r['unified_label']] += 1
	labels = sorted([l for l, n in cls_count.items() if n >= args.min_samples])
	label_to_idx = {l: i for i, l in enumerate(labels)}
	n_cls = len(label_to_idx)
	log(f"Classes (min_samples={args.min_samples}): {n_cls}")

	# Build frozen backbone (ImageNet/SSL pretrained, NO classification head training)
	model = build_backbone(args.backbone_arch, ssl_ckpt=args.ssl_ckpt, n_classes=n_cls).to(device)
	# Freeze everything
	for p in model.parameters(): p.requires_grad = False

	from torchvision import transforms
	img_size = get_arch_img_size(args.backbone_arch)
	tf = transforms.Compose([
	transforms.Resize((img_size, img_size), antialias=True),
	transforms.ToTensor(),
	transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
	])
	tr_ds = SimpleDS(args.manifest, label_to_idx, tf, val_fold=args.val_fold, is_val=False,
	min_samples=args.min_samples)
	va_ds = SimpleDS(args.manifest, label_to_idx, tf, val_fold=args.val_fold, is_val=True,
	min_samples=args.min_samples)
	log(f"Train={len(tr_ds)} Val={len(va_ds)}")

	tr_dl = DataLoader(tr_ds, batch_size=args.batch_size, shuffle=False, num_workers=6)
	va_dl = DataLoader(va_ds, batch_size=args.batch_size, shuffle=False, num_workers=4)

	log("Extracting train features (frozen)...")
	tr_f, tr_y = extract(model, tr_dl, device, dtype)
	log("Extracting val features...")
	va_f, va_y = extract(model, va_dl, device, dtype)
	D = tr_f.size(-1)
	log(f"Feature dim: {D}")

	# Linear probe on extracted feats (tiny: just one Linear layer)
	head = nn.Linear(D, n_cls).to(device)
	opt = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=args.weight_decay)
	sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)

	# Batch over features in memory
	tr_f_d, tr_y_d = tr_f.to(device), tr_y.to(device)
	va_f_d, va_y_d = va_f.to(device), va_y.to(device)
	best_acc = 0.0
	for ep in range(args.epochs):
	head.train()
	idx = torch.randperm(len(tr_f_d), device=device)
	for i in range(0, len(tr_f_d), 1024):
	b = idx[i:i+1024]
	logits = head(tr_f_d[b])
	loss = F.cross_entropy(logits, tr_y_d[b], label_smoothing=0.1)
	opt.zero_grad(); loss.backward(); opt.step()
	sched.step()
	# Validate
	head.eval()
	with torch.no_grad():
	logits = head(va_f_d)
	acc = (logits.argmax(-1) == va_y_d).float().mean().item()
	_, top5 = logits.topk(5, dim=-1)
	top5_acc = sum(va_y_d[i].item() in top5[i].tolist()
	for i in range(len(va_y_d))) / len(va_y_d)
	if ep % 5 == 0 or ep == args.epochs - 1:
	log(f"ep {ep}: val_top1={acc:.4f} val_top5={top5_acc:.4f} loss={loss.item():.3f}")
	if acc > best_acc:
	best_acc = acc
	best_probs = F.softmax(head(va_f_d).float(), dim=-1).cpu()

	log(f"=== Frozen {args.backbone_arch} linear probe ===")
	log(f" Best top-1: {best_acc:.4f}")

	Path(args.output).parent.mkdir(parents=True, exist_ok=True)
	torch.save({
	'probs': best_probs, 'targets': va_y,
	'top1': best_acc, 'label_to_idx': label_to_idx,
	'backbone': args.backbone_arch, 'feat_dim': D,
	'head_state': head.state_dict(),
	}, args.output)
	log(f"Saved → {args.output}")

	if __name__ == '__main__':
	main()