""" finetune_backbone.py Fine-tune ResNet18 (ImageNet) on the local charcoal microscopy dataset. Goal: produce a domain-adapted backbone for students to use as a frozen feature extractor. The full dataset is used intentionally — this is a teaching artifact, not a research model with a held-out test split. Output (in backbone/): resnet18_charcoal_backbone.pt — backbone weights, FC replaced by Identity backbone_meta.json — class names, feature dim, training info Usage: python finetune_backbone.py python finetune_backbone.py --epochs 40 --batch-size 16 """ import argparse import json import time from pathlib import Path import torch import torch.nn as nn import torch.optim as optim from PIL import Image from torch.utils.data import DataLoader, Dataset from torchvision import models, transforms # --------------------------------------------------------------------------- # Paths # --------------------------------------------------------------------------- ROOT = Path(__file__).parent DATA_DIR = ROOT / "data" OUTPUT_DIR = ROOT / "backbone" OUTPUT_DIR.mkdir(exist_ok=True) # --------------------------------------------------------------------------- # Defaults # --------------------------------------------------------------------------- IMAGE_SIZE = 224 SEED = 42 WARMUP_EPOCHS = 10 # backbone frozen, only FC trained WARMUP_LR = 1e-3 FINETUNE_EPOCHS = 40 # all layers unfrozen, small LR FINETUNE_LR = 5e-5 WEIGHT_DECAY = 1e-4 # --------------------------------------------------------------------------- # Dataset # --------------------------------------------------------------------------- class CharcoalDataset(Dataset): """Flat ImageFolder-style dataset that handles .tif files.""" EXTENSIONS = {".tif", ".tiff", ".jpg", ".jpeg", ".png"} def __init__(self, root: Path, transform=None): self.transform = transform self.classes = sorted( d.name for d in root.iterdir() if d.is_dir() and not d.name.startswith(".") ) self.class_to_idx = {c: i for i, c in enumerate(self.classes)} self.samples = [] for cls in self.classes: for p in sorted((root / cls).iterdir()): if p.suffix.lower() in self.EXTENSIONS: self.samples.append((p, self.class_to_idx[cls])) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label = self.samples[idx] image = Image.open(path).convert("RGB") if self.transform: image = self.transform(image) return image, label def make_transform(): # Aggressive augmentation: microscopy images have no canonical orientation # and vary in staining intensity. return transforms.Compose([ transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.RandomRotation(180), transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2), transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.85, 1.15)), transforms.ToTensor(), transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ]) # --------------------------------------------------------------------------- # Training helpers # --------------------------------------------------------------------------- def run_epoch(model, loader, criterion, optimizer, device): model.train() total_loss, correct, total = 0.0, 0, 0 for images, labels in loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() total_loss += loss.item() * images.size(0) correct += (outputs.argmax(1) == labels).sum().item() total += labels.size(0) return total_loss / total, correct / total # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser() parser.add_argument("--warmup-epochs", type=int, default=WARMUP_EPOCHS) parser.add_argument("--finetune-epochs", type=int, default=FINETUNE_EPOCHS) parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--warmup-lr", type=float, default=WARMUP_LR) parser.add_argument("--finetune-lr", type=float, default=FINETUNE_LR) args = parser.parse_args() torch.manual_seed(SEED) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device : {device}") dataset = CharcoalDataset(DATA_DIR, transform=make_transform()) num_classes = len(dataset.classes) print(f"Classes : {num_classes} | Images : {len(dataset)}") print(f" {', '.join(dataset.classes)}\n") loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=True, num_workers=0, # 0 = safe on Windows pin_memory=(device.type == "cuda"), ) # ----------------------------------------------------------------------- # Build model # ----------------------------------------------------------------------- model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT) model.fc = nn.Linear(model.fc.in_features, num_classes) model.to(device) # Label smoothing helps regularise with tiny datasets criterion = nn.CrossEntropyLoss(label_smoothing=0.1) # ----------------------------------------------------------------------- # Phase 1 — warm-up: freeze backbone, train FC only # ----------------------------------------------------------------------- print(f"=== Phase 1 : warm-up ({args.warmup_epochs} epochs, backbone frozen) ===") for p in model.parameters(): p.requires_grad = False for p in model.fc.parameters(): p.requires_grad = True optimizer = optim.AdamW(model.fc.parameters(), lr=args.warmup_lr, weight_decay=WEIGHT_DECAY) for epoch in range(1, args.warmup_epochs + 1): loss, acc = run_epoch(model, loader, criterion, optimizer, device) print(f" [{epoch:>3}/{args.warmup_epochs}] loss={loss:.4f} acc={acc:.4f}") # ----------------------------------------------------------------------- # Phase 2 — full fine-tune: unfreeze all layers # ----------------------------------------------------------------------- print(f"\n=== Phase 2 : fine-tune ({args.finetune_epochs} epochs, all layers) ===") for p in model.parameters(): p.requires_grad = True optimizer = optim.AdamW( model.parameters(), lr=args.finetune_lr, weight_decay=WEIGHT_DECAY ) scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.finetune_epochs, eta_min=args.finetune_lr * 0.05 ) best_acc = 0.0 best_state = None t0 = time.time() for epoch in range(1, args.finetune_epochs + 1): loss, acc = run_epoch(model, loader, criterion, optimizer, device) scheduler.step() lr = optimizer.param_groups[0]["lr"] print(f" [{epoch:>3}/{args.finetune_epochs}] loss={loss:.4f} acc={acc:.4f} lr={lr:.2e}") if acc > best_acc: best_acc = acc best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} elapsed = time.time() - t0 print(f"\nTemps phase 2 : {elapsed:.0f}s | Meilleure accuracy entraînement : {best_acc:.4f}") # ----------------------------------------------------------------------- # Save backbone (FC replaced by Identity — outputs 512-dim feature vector) # ----------------------------------------------------------------------- model.load_state_dict(best_state) backbone = models.resnet18() backbone.fc = nn.Identity() # Transfer all weights except fc (which is now Identity with no parameters) backbone_state = {k: v for k, v in best_state.items() if not k.startswith("fc.")} backbone.load_state_dict(backbone_state, strict=False) backbone_path = OUTPUT_DIR / "resnet18_charcoal_backbone.pt" torch.save(backbone.state_dict(), backbone_path) print(f"Backbone sauvegardé : {backbone_path}") # ----------------------------------------------------------------------- # Save metadata # ----------------------------------------------------------------------- meta = { "classes": dataset.classes, "num_classes": num_classes, "image_size": IMAGE_SIZE, "feature_dim": 512, "warmup_epochs": args.warmup_epochs, "finetune_epochs": args.finetune_epochs, "best_train_acc": round(float(best_acc), 4), "device": str(device), } meta_path = OUTPUT_DIR / "backbone_meta.json" with open(meta_path, "w", encoding="utf-8") as f: json.dump(meta, f, indent=2, ensure_ascii=False) print(f"Métadonnées sauvegardées : {meta_path}") if __name__ == "__main__": main()