""" train_clf.py — InceptionV3 binary classification on BDAPPV (panel / no panel). Usage: python train_clf.py --provider google python train_clf.py --provider ign python train_clf.py --provider google --checkpoint weights/inception_google_best.pth """ import argparse import os import torch from tqdm import tqdm import torch.nn as nn from torch.utils.data import DataLoader from torchvision.models import inception_v3, Inception_V3_Weights from datasets import load_dataset from dataset import ClassificationDataset # ── Metrics ─────────────────────────────────────────────────────────────────── def accuracy_f1(logits, labels, threshold=0.0): preds = (logits > threshold).float() correct = (preds == labels).float().mean().item() tp = (preds * labels).sum().item() fp = (preds * (1 - labels)).sum().item() fn = ((1 - preds) * labels).sum().item() prec = tp / max(tp + fp, 1e-6) rec = tp / max(tp + fn, 1e-6) f1 = 2 * prec * rec / max(prec + rec, 1e-6) return correct, f1 # ── Main ────────────────────────────────────────────────────────────────────── def train(args): device = ( torch.device("mps") if torch.backends.mps.is_available() else torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") ) print(f"Device : {device}") print(f"Provider: {args.provider} | effective batch: {args.batch_size * args.grad_accum}") # ── Data ────────────────────────────────────────────────────────────────── print("Loading dataset from HuggingFace Hub...") ds = load_dataset("gabrielkasmi/bdappv", args.provider) train_ds = ClassificationDataset(ds["train"], img_size=299, augment=True) val_ds = ClassificationDataset(ds["validation"], img_size=299, augment=False) print(f" train : {len(train_ds):,} | val : {len(val_ds):,}") train_loader = DataLoader( train_ds, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=(device.type != "mps"), ) val_loader = DataLoader( val_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=(device.type != "mps"), ) # ── Model ───────────────────────────────────────────────────────────────── model = inception_v3(weights=Inception_V3_Weights.DEFAULT) if args.checkpoint: # Poids issus d'un modèle à 2 classes — on charge avec la tête originale # puis on la remplace par une tête binaire (1 sortie, BCEWithLogitsLoss) model.fc = nn.Linear(model.fc.in_features, 2) model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, 2) ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) state = ckpt["model_state_dict"] if isinstance(ckpt, dict) else ckpt model.load_state_dict(state, strict=True) print(f"Loaded checkpoint: {args.checkpoint}") # Remplace la tête par 1 sortie pour BCEWithLogitsLoss model.fc = nn.Linear(model.fc.in_features, 1) model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, 1) model = model.to(device) # ── Optimizer + scheduler ───────────────────────────────────────────────── optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) criterion = nn.BCEWithLogitsLoss() # ── Training loop ───────────────────────────────────────────────────────── os.makedirs(args.out_dir, exist_ok=True) best_f1 = 0.0 patience_ctr = 0 for epoch in range(1, args.epochs + 1): model.train() total_loss = 0.0 optimizer.zero_grad() pbar = tqdm(train_loader, desc=f"Epoch {epoch}/{args.epochs}", leave=False) for step, (images, labels) in enumerate(pbar): images = images.to(device) labels = labels.to(device).unsqueeze(1) out = model(images) logits = out.logits if hasattr(out, "logits") else out[0] aux = out.aux_logits if hasattr(out, "aux_logits") else out[1] loss = (criterion(logits, labels) + 0.4 * criterion(aux, labels)) / args.grad_accum loss.backward() total_loss += loss.item() * args.grad_accum pbar.set_postfix(loss=f"{loss.item() * args.grad_accum:.4f}") if (step + 1) % args.grad_accum == 0 or (step + 1) == len(train_loader): optimizer.step() optimizer.zero_grad() # ── Validation ──────────────────────────────────────────────────────── model.eval() all_logits, all_labels = [], [] with torch.no_grad(): for images, labels in val_loader: images = images.to(device) logits = model(images) if hasattr(logits, "logits"): logits = logits.logits all_logits.append(logits.cpu()) all_labels.append(labels.unsqueeze(1)) all_logits = torch.cat(all_logits) all_labels = torch.cat(all_labels) acc, f1 = accuracy_f1(all_logits, all_labels) avg_loss = total_loss / len(train_loader) print(f"Epoch {epoch:3d}/{args.epochs} loss={avg_loss:.4f} val_acc={acc:.4f} val_f1={f1:.4f}") scheduler.step() if f1 > best_f1: best_f1 = f1 patience_ctr = 0 ckpt_path = os.path.join(args.out_dir, f"inception_{args.provider}_best.pth") torch.save(model.state_dict(), ckpt_path) print(f" → checkpoint saved (F1={best_f1:.4f})") else: patience_ctr += 1 if patience_ctr >= args.patience: print(f"\nEarly stopping — no improvement for {args.patience} epochs.") break print(f"\nDone. Best val F1: {best_f1:.4f}") print(f"Weights: {os.path.join(args.out_dir, f'inception_{args.provider}_best.pth')}") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--provider", default="google", choices=["google", "ign"]) parser.add_argument("--epochs", type=int, default=50) parser.add_argument("--batch_size", type=int, default=8) parser.add_argument("--grad_accum", type=int, default=4, help="Gradient accumulation steps") parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--num_workers", type=int, default=2, help="Workers pour le dataloader") parser.add_argument("--checkpoint", default=None) parser.add_argument("--out_dir", default="weights") parser.add_argument("--patience", type=int, default=7, help="Early stopping patience (epochs)") args = parser.parse_args() train(args)