From 9dd8b3007e79f0634f868d405f5e7c48917305fd Mon Sep 17 00:00:00 2001 From: Vineet Kumar Date: Mon, 20 Apr 2026 10:03:17 -0400 Subject: improve training for class-imbalanced DR classification - WeightedRandomSampler for balanced class sampling during training - FocalLoss (gamma=2) for hard-example mining, no alpha (sampler handles balance) - Stronger augmentation: RandomResizedCrop, ColorJitter, GaussianBlur, 30deg rotation - Discriminative LR: backbone at 0.1x, head at 1x - 3-epoch linear warmup + cosine decay scheduler - Early stopping on macro F1 instead of weighted F1 - Updated defaults: batch_size=64, epochs=50, patience=10 - Comparison chart now shows both macro and weighted F1 - DataParallel checkpoint handling (strip module. prefix) - Enable cudnn.benchmark for training speed --- main.py | 60 +++++++++++++++++++++++++++++++++++++++++++----------------- 1 file changed, 43 insertions(+), 17 deletions(-) (limited to 'main.py') diff --git a/main.py b/main.py index 6d8b2d9..fb87c33 100644 --- a/main.py +++ b/main.py @@ -12,10 +12,10 @@ import os import torch import torch.nn as nn -from torch.optim.lr_scheduler import CosineAnnealingLR +from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR import kagglehub -from utils import set_seed, compute_class_weights, get_device, CLASS_NAMES +from utils import set_seed, get_device, FocalLoss from dataset import split_dataset, build_dataloaders from models import create_model from train import train_model @@ -35,15 +35,17 @@ def parse_args(): '--clahe', nargs='+', type=int, default=[0, 1], help='CLAHE settings to run (0=off, 1=on)' ) - parser.add_argument('--epochs', type=int, default=30, help='Max training epochs') - parser.add_argument('--patience', type=int, default=5, help='Early stopping patience') - parser.add_argument('--batch-size', type=int, default=32, help='Training batch size') + parser.add_argument('--epochs', type=int, default=50, help='Max training epochs') + parser.add_argument('--patience', type=int, default=10, help='Early stopping patience') + parser.add_argument('--batch-size', type=int, default=64, help='Training batch size') parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate') parser.add_argument('--weight-decay', type=float, default=1e-4, help='AdamW weight decay') parser.add_argument('--seed', type=int, default=42, help='Random seed') parser.add_argument('--data-root', type=str, default=None, help='Path to dataset root (auto-downloaded if not set)') - parser.add_argument('--num-workers', type=int, default=4, help='DataLoader workers') + parser.add_argument('--num-workers', type=int, default=8, help='DataLoader workers') + parser.add_argument('--focal-gamma', type=float, default=2.0, help='Focal loss gamma (0=standard CE)') + parser.add_argument('--warmup-epochs', type=int, default=3, help='LR warmup epochs') return parser.parse_args() @@ -64,10 +66,6 @@ def main(): train_paths, train_labels, val_paths, val_labels, test_paths, test_labels = \ split_dataset(data_root, seed=args.seed) - # Class weights computed once from training labels - class_weights = compute_class_weights(train_labels).to(device) - print("Class weights:", [f"{w:.3f}" for w in class_weights.cpu().numpy()]) - # ── Experiments ────────────────────────────────────────────────────────── experiments = [ (model_name, bool(use_clahe)) @@ -97,11 +95,35 @@ def main(): ) model = create_model(model_name, num_classes=5, pretrained=True).to(device) - criterion = nn.CrossEntropyLoss(weight=class_weights) + + # Discriminative LR: backbone gets 10x lower LR than the new head + if model_name == 'resnet50': + head_params = list(model.fc.parameters()) + elif model_name == 'efficientnet_b0': + head_params = list(model.classifier.parameters()) + elif model_name == 'vit_b_16': + head_params = list(model.heads.parameters()) + else: + head_params = [] + head_ids = {id(p) for p in head_params} + backbone_params = [p for p in model.parameters() if id(p) not in head_ids] + + if torch.cuda.device_count() > 1: + print(f" Using {torch.cuda.device_count()} GPUs via DataParallel") + model = torch.nn.DataParallel(model) + + criterion = FocalLoss(alpha=None, gamma=args.focal_gamma) optimizer = torch.optim.AdamW( - model.parameters(), lr=args.lr, weight_decay=args.weight_decay + [ + {'params': backbone_params, 'lr': args.lr * 0.1}, + {'params': head_params, 'lr': args.lr}, + ], + weight_decay=args.weight_decay, ) - scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs) + warmup_epochs = min(args.warmup_epochs, args.epochs - 1) + warmup_scheduler = LinearLR(optimizer, start_factor=0.1, total_iters=warmup_epochs) + cosine_scheduler = CosineAnnealingLR(optimizer, T_max=max(args.epochs - warmup_epochs, 1)) + scheduler = SequentialLR(optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_epochs]) # Train train_results = train_model( @@ -113,16 +135,20 @@ def main(): model_name=exp_key, ) - # Save best checkpoint + # Save best checkpoint (strip DataParallel 'module.' prefix for portability) ckpt_path = os.path.join(exp_dir, 'best_model.pth') - torch.save(train_results['best_model_state'], ckpt_path) + best_state = train_results['best_model_state'] + if isinstance(model, torch.nn.DataParallel): + best_state = {k[7:]: v for k, v in best_state.items()} + torch.save(best_state, ckpt_path) print(f" Checkpoint saved: {ckpt_path}") # Load best weights for evaluation - model.load_state_dict(train_results['best_model_state']) + core_model = model.module if isinstance(model, torch.nn.DataParallel) else model + core_model.load_state_dict(best_state) # Evaluate on test set - test_metrics = evaluate_model(model, test_loader, device) + test_metrics = evaluate_model(core_model, test_loader, device) print(f"\n Test results — Weighted F1: {test_metrics['weighted_f1']:.4f} | " f"Macro F1: {test_metrics['macro_f1']:.4f} | " -- cgit v1.2.3