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, 28 insertions, 2 deletions

diff --git a/utils.py b/utils.py
index a836ca2..90c79d5 100644
--- a/utils.py
+++ b/utils.py
@@ -1,6 +1,8 @@
 import random
 import numpy as np
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 
 CLASS_NAMES = [
     'Healthy (0)',
@@ -18,8 +20,7 @@ def set_seed(seed: int = 42) -> None:
     np.random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.benchmark = True
 
 
 def compute_class_weights(labels) -> torch.FloatTensor:
@@ -35,3 +36,28 @@ def get_device() -> torch.device:
     if torch.cuda.is_available():
         return torch.device('cuda')
     return torch.device('cpu')
+
+
+class FocalLoss(nn.Module):
+    """Focal loss with optional per-class alpha weighting.
+
+    Focuses learning on hard, misclassified examples by down-weighting
+    easy examples (high pt). Good for class imbalance in medical imaging.
+
+    Args:
+        alpha: Per-class weight tensor (same shape as class weights for CE).
+               If None, no per-class weighting is applied.
+        gamma: Focusing parameter. gamma=0 reduces to standard CE.
+               gamma=2 is the standard value from the RetinaNet paper.
+    """
+
+    def __init__(self, alpha=None, gamma: float = 2.0):
+        super().__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+
+    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        ce_loss = F.cross_entropy(inputs, targets, weight=self.alpha, reduction='none')
+        pt = torch.exp(-ce_loss)
+        focal_loss = ((1.0 - pt) ** self.gamma) * ce_loss
+        return focal_loss.mean()