| """ |
| 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 |
|
|
|
|
| |
|
|
| 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 |
|
|
|
|
| |
|
|
| 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}") |
|
|
| |
| 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 = inception_v3(weights=Inception_V3_Weights.DEFAULT) |
|
|
| if args.checkpoint: |
| |
| |
| 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}") |
|
|
| |
| 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 = 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() |
|
|
| |
| 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() |
|
|
| |
| 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) |
|
|
