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

1 files changed, 43 insertions, 17 deletions

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} | "