"""Fine-tune DINOv3 backbone + classification head trên dữ liệu cân bằng. Thiết kế: - Backbone: DINOv3 ViT-S/16 (load từ model.safetensors) — được fine-tune (không đông cứng) - Head: Linear(384, 2) trên CLS token - Optimizer: AdamW 2 nhóm LR (backbone thấp 1e-5, head cao 1e-3) - Scheduler: CosineAnnealingLR - Lưu checkpoint val acc tốt nhất; sau cùng đánh giá trên test Cách chạy: .venv/bin/python scripts/train.py --train-csv data/splits/train_insight.csv \ --val-csv data/splits/val_insight.csv --test-csv data/splits/test_insight.csv """ import argparse import csv import json import os import sys import time import numpy as np import torch import torch.nn as nn from PIL import Image from sklearn.metrics import ( accuracy_score, confusion_matrix, f1_score, precision_score, recall_score, roc_auc_score, ) from torch.utils.data import DataLoader, Dataset from torchvision import transforms from tqdm import tqdm sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) from src.models.dinov3_vit import load_dinov3 # noqa: E402 IMG_SIZE = 256 MEAN, STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] TRAIN_TF = transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(MEAN, STD), ]) EVAL_TF = transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize(MEAN, STD), ]) class ImageDataset(Dataset): def __init__(self, csv_path, transform=None): self.rows = [] with open(csv_path, newline="") as f: reader = csv.reader(f) next(reader, None) for row in reader: if len(row) >= 2: self.rows.append((row[0], int(row[1]))) self.transform = transform def __len__(self): return len(self.rows) def __getitem__(self, i): path, label = self.rows[i] img = Image.open(path).convert("RGB") if self.transform: img = self.transform(img) return img, label class DinoViTClassifier(nn.Module): """Backbone DINOv3 + Linear head trên CLS token.""" def __init__(self, backbone: nn.Module, num_classes: int = 2): super().__init__() self.backbone = backbone self.head = nn.Linear(backbone.embed_dim, num_classes) def forward(self, x): return self.head(self.backbone(x)) @torch.no_grad() def evaluate(model, loader, device): model.eval() all_y, all_pred, all_prob = [], [], [] for x, y in loader: x = x.to(device) logits = model(x) probs = torch.softmax(logits, dim=1) all_y.extend(y.tolist()) all_pred.extend(logits.argmax(1).tolist()) all_prob.extend(probs[:, 1].tolist()) y = np.array(all_y) pred = np.array(all_pred) prob = np.array(all_prob) return { "accuracy": float(accuracy_score(y, pred)), "precision": float(precision_score(y, pred)), "recall": float(recall_score(y, pred)), "f1": float(f1_score(y, pred)), "roc_auc": float(roc_auc_score(y, prob)), "confusion_matrix": confusion_matrix(y, pred, labels=[0, 1]).tolist(), } def main(): parser = argparse.ArgumentParser(description="Fine-tune DINOv3") parser.add_argument("--train-csv", required=True) parser.add_argument("--val-csv", required=True) parser.add_argument("--test-csv", required=True) parser.add_argument("--model", default="models/model.safetensors") parser.add_argument("--output-dir", default="outputs/checkpoints") parser.add_argument("--report", default="outputs/results/finetune_report.json") parser.add_argument("--epochs", type=int, default=5) parser.add_argument("--batch-size", type=int, default=32) parser.add_argument("--lr-backbone", type=float, default=1e-5) parser.add_argument("--lr-head", type=float, default=1e-3) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--device", default="auto", choices=["auto", "cuda", "mps", "cpu"]) args = parser.parse_args() if args.device == "auto": device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" else: device = args.device print(f"Device: {device}", flush=True) torch.manual_seed(args.seed) np.random.seed(args.seed) # ---------- Data ---------- train_ds = ImageDataset(args.train_csv, TRAIN_TF) val_ds = ImageDataset(args.val_csv, EVAL_TF) test_ds = ImageDataset(args.test_csv, EVAL_TF) train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=0) val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=0) test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False, num_workers=0) print(f"Train: {len(train_ds)} | Val: {len(val_ds)} | Test: {len(test_ds)}", flush=True) # ---------- Model ---------- backbone = load_dinov3(args.model, img_size=IMG_SIZE) model = DinoViTClassifier(backbone).to(device) n_params = sum(p.numel() for p in model.parameters()) print(f"Model: {n_params:,} params ({n_params/1e6:.1f}M) — sẽ fine-tune CẢ backbone", flush=True) # ---------- Optimizer / Scheduler ---------- optimizer = torch.optim.AdamW([ {"params": model.backbone.parameters(), "lr": args.lr_backbone}, {"params": model.head.parameters(), "lr": args.lr_head}, ], weight_decay=0.05) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) criterion = nn.CrossEntropyLoss() # ---------- Train ---------- os.makedirs(args.output_dir, exist_ok=True) best_path = os.path.join(args.output_dir, "dinov3_finetuned.pt") best_val_acc = 0.0 history = [] for epoch in range(args.epochs): model.train() t0 = time.time() total_loss, n_batch = 0.0, 0 pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}", unit="batch", ncols=100, leave=True) for x, y in pbar: x, y = x.to(device), y.to(device) optimizer.zero_grad() out = model(x) loss = criterion(out, y) loss.backward() optimizer.step() total_loss += loss.item() n_batch += 1 if n_batch % 20 == 0: pbar.set_postfix(loss=f"{loss.item():.4f}") pbar.close() scheduler.step() val_metrics = evaluate(model, val_loader, device) train_loss = total_loss / n_batch history.append({"epoch": epoch + 1, "train_loss": train_loss, **val_metrics}) dt = time.time() - t0 print(f"[Epoch {epoch+1}/{args.epochs}] loss={train_loss:.4f} | " f"val_acc={val_metrics['accuracy']:.4f} val_f1={val_metrics['f1']:.4f} | {dt:.0f}s", flush=True) if val_metrics["accuracy"] > best_val_acc: best_val_acc = val_metrics["accuracy"] torch.save({"state_dict": model.state_dict(), "epoch": epoch + 1, "val_metrics": val_metrics}, best_path) print(f" → Lưu checkpoint tốt nhất (val_acc={best_val_acc:.4f})", flush=True) # ---------- Test ---------- print(f"\nLoad checkpoint tốt nhất và đánh giá TEST...") ckpt = torch.load(best_path, map_location=device) model.load_state_dict(ckpt["state_dict"]) test_metrics = evaluate(model, test_loader, device) print("\n================= KẾT QUẢ TEST (sau fine-tune) =================") print(f" accuracy : {test_metrics['accuracy']:.4f}") print(f" precision: {test_metrics['precision']:.4f}") print(f" recall : {test_metrics['recall']:.4f}") print(f" f1 : {test_metrics['f1']:.4f}") print(f" roc_auc : {test_metrics['roc_auc']:.4f}") cm = test_metrics["confusion_matrix"] print(f" CM [TN FP; FN TP]: TN={cm[0][0]} FP={cm[0][1]} FN={cm[1][0]} TP={cm[1][1]}") print("===================================================================") report = {"best_val_acc": best_val_acc, "test": test_metrics, "history": history, "args": vars(args)} os.makedirs(os.path.dirname(args.report), exist_ok=True) with open(args.report, "w") as f: json.dump(report, f, indent=2, ensure_ascii=False) print(f"Đã lưu: {best_path} + {args.report}") if __name__ == "__main__": main()