commit 9dd8b3007e79f0634f868d405f5e7c48917305fd
parent 8f571ee51854540d186aafc01d25908bba5ea09e
Author: Vineet Kumar <git@vineetk.net>
Date: Mon, 20 Apr 2026 10:03:17 -0400
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
Diffstat:
| M | dataset.py | | | 70 | +++++++++++++++++++++++++++++++++++++++++++++++++++------------------- |
| M | evaluate.py | | | 34 | ++++++++++++++++++++++------------ |
| M | main.py | | | 60 | +++++++++++++++++++++++++++++++++++++++++++----------------- |
| M | train.py | | | 29 | ++++++++++++++++------------- |
| M | utils.py | | | 30 | ++++++++++++++++++++++++++++-- |
5 files changed, 160 insertions(+), 63 deletions(-)
diff --git a/dataset.py b/dataset.py
@@ -4,7 +4,7 @@ import numpy as np
import cv2
import torch
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from torchvision import transforms
from sklearn.model_selection import train_test_split
@@ -19,14 +19,12 @@ class CLAHETransform:
"""
def __init__(self, clip_limit: float = 2.0, tile_grid_size: tuple = (8, 8)):
- self.clip_limit = clip_limit
- self.tile_grid_size = tile_grid_size
+ self.clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
def __call__(self, img: Image.Image) -> Image.Image:
img_np = np.array(img) # RGB uint8
lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
- clahe = cv2.createCLAHE(clipLimit=self.clip_limit, tileGridSize=self.tile_grid_size)
- lab[:, :, 0] = clahe.apply(lab[:, :, 0])
+ lab[:, :, 0] = self.clahe.apply(lab[:, :, 0])
result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
return Image.fromarray(result)
@@ -77,17 +75,32 @@ def split_dataset(data_root: str, seed: int = 42):
return train_paths, train_labels, val_paths, val_labels, test_paths, test_labels
-def build_transforms(use_clahe: bool, training: bool):
- """Build a transform pipeline for training or eval."""
- ops = [transforms.Resize((224, 224))]
- if use_clahe:
- ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8)))
+def build_transforms(use_clahe: bool, training: bool, img_size: int = 224):
+ """Build a transform pipeline for training or eval.
+
+ Training path uses RandomResizedCrop for scale/position augmentation.
+ Eval path uses a deterministic Resize to preserve comparability.
+ CLAHE is applied before spatial transforms at a slightly larger size
+ so the crop has room to operate.
+ """
+ ops = []
if training:
+ if use_clahe:
+ # Resize larger so RandomResizedCrop still sees enough context after CLAHE
+ ops.append(transforms.Resize((img_size + 32, img_size + 32)))
+ ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8)))
ops += [
+ transforms.RandomResizedCrop(img_size, scale=(0.8, 1.0), ratio=(0.9, 1.1)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomVerticalFlip(p=0.5),
- transforms.RandomRotation(degrees=15),
+ transforms.RandomRotation(degrees=30),
+ transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.02),
+ transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 1.0)),
]
+ else:
+ ops.append(transforms.Resize((img_size, img_size)))
+ if use_clahe:
+ ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8)))
ops += [
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
@@ -102,22 +115,41 @@ def build_dataloaders(
use_clahe: bool,
batch_size: int = 32,
num_workers: int = 4,
+ img_size: int = 224,
):
- """Create DataLoaders for all three splits."""
- train_ds = DRDataset(train_paths, train_labels, transform=build_transforms(use_clahe, training=True))
- val_ds = DRDataset(val_paths, val_labels, transform=build_transforms(use_clahe, training=False))
- test_ds = DRDataset(test_paths, test_labels, transform=build_transforms(use_clahe, training=False))
+ """Create DataLoaders for all three splits.
+
+ The training DataLoader uses WeightedRandomSampler so that each class
+ is sampled at approximately equal frequency, counteracting class imbalance.
+ """
+ train_ds = DRDataset(train_paths, train_labels, transform=build_transforms(use_clahe, training=True, img_size=img_size))
+ val_ds = DRDataset(val_paths, val_labels, transform=build_transforms(use_clahe, training=False, img_size=img_size))
+ test_ds = DRDataset(test_paths, test_labels, transform=build_transforms(use_clahe, training=False, img_size=img_size))
+
+ # Per-sample weights: inverse of class frequency so all classes are seen equally
+ labels_arr = np.array(train_labels)
+ class_counts = np.bincount(labels_arr, minlength=5)
+ class_sample_weights = 1.0 / class_counts.astype(float)
+ sample_weights = class_sample_weights[labels_arr]
+ sampler = WeightedRandomSampler(
+ weights=torch.from_numpy(sample_weights).float(),
+ num_samples=len(train_labels),
+ replacement=True,
+ )
train_loader = DataLoader(
- train_ds, batch_size=batch_size, shuffle=True,
- num_workers=num_workers, pin_memory=True
+ train_ds, batch_size=batch_size, sampler=sampler,
+ num_workers=num_workers, pin_memory=True,
+ persistent_workers=True, prefetch_factor=4,
)
val_loader = DataLoader(
val_ds, batch_size=batch_size * 2, shuffle=False,
- num_workers=num_workers, pin_memory=True
+ num_workers=num_workers, pin_memory=True,
+ persistent_workers=True, prefetch_factor=4,
)
test_loader = DataLoader(
test_ds, batch_size=batch_size * 2, shuffle=False,
- num_workers=num_workers, pin_memory=True
+ num_workers=num_workers, pin_memory=True,
+ persistent_workers=True, prefetch_factor=4,
)
return train_loader, val_loader, test_loader
diff --git a/evaluate.py b/evaluate.py
@@ -76,10 +76,10 @@ def plot_training_curves(
ax1.set_title('Training and Validation Loss')
ax1.legend()
- ax2.plot(epochs, val_f1s, label='Val weighted F1', color='green')
+ ax2.plot(epochs, val_f1s, label='Val macro F1', color='green')
ax2.set_xlabel('Epoch')
- ax2.set_ylabel('Weighted F1')
- ax2.set_title('Validation Weighted F1')
+ ax2.set_ylabel('Macro F1')
+ ax2.set_title('Validation Macro F1')
ax2.legend()
plt.tight_layout()
@@ -127,18 +127,28 @@ def generate_summary(all_results: dict, save_dir: str) -> None:
writer.writerows(rows)
print(f"\nSummary saved to {csv_path}")
- # Comparison bar chart
- labels = [f"{r['model']}\nCLAHE={r['clahe']}" for r in rows]
+ # Comparison bar chart — show both macro and weighted F1
+ x_labels = [f"{r['model']}\nCLAHE={r['clahe']}" for r in rows]
w_f1_vals = [r['weighted_f1'] for r in rows]
-
- fig, ax = plt.subplots(figsize=(10, 5))
- bars = ax.bar(labels, w_f1_vals, color=['#4C72B0', '#DD8452'] * 3)
- ax.set_ylabel('Weighted F1-Score')
- ax.set_title('Weighted F1 Comparison Across Experiments')
+ m_f1_vals = [r['macro_f1'] for r in rows]
+
+ x = np.arange(len(x_labels))
+ bar_width = 0.35
+ fig, ax = plt.subplots(figsize=(12, 5))
+ bars_w = ax.bar(x - bar_width / 2, w_f1_vals, bar_width, label='Weighted F1', color='#4C72B0')
+ bars_m = ax.bar(x + bar_width / 2, m_f1_vals, bar_width, label='Macro F1', color='#DD8452')
+ ax.set_ylabel('F1-Score')
+ ax.set_title('F1 Comparison Across Experiments')
+ ax.set_xticks(x)
+ ax.set_xticklabels(x_labels)
ax.set_ylim(0, 1.0)
- for bar, val in zip(bars, w_f1_vals):
+ ax.legend()
+ for bar, val in zip(bars_w, w_f1_vals):
+ ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
+ f'{val:.3f}', ha='center', va='bottom', fontsize=8)
+ for bar, val in zip(bars_m, m_f1_vals):
ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
- f'{val:.4f}', ha='center', va='bottom', fontsize=9)
+ f'{val:.3f}', ha='center', va='bottom', fontsize=8)
plt.tight_layout()
chart_path = os.path.join(save_dir, 'comparison_chart.png')
fig.savefig(chart_path, dpi=150)
diff --git 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} | "
diff --git a/train.py b/train.py
@@ -19,11 +19,11 @@ def train_model(
patience: int = 5,
model_name: str = "",
) -> dict:
- """Train model with mixed precision, early stopping on val weighted-F1.
+ """Train model with mixed precision, early stopping on val macro F1.
Returns a dict with:
- best_model_state, train_losses, val_losses, val_f1s,
- train_time (seconds), epochs_trained
+ best_model_state, train_losses, val_losses, val_f1s (macro), val_weighted_f1s,
+ train_time (seconds), epochs_trained, best_val_f1 (macro)
"""
scaler = torch.amp.GradScaler('cuda') if device.type == 'cuda' else None
@@ -31,7 +31,7 @@ def train_model(
best_state = None
patience_counter = 0
- train_losses, val_losses, val_f1s = [], [], []
+ train_losses, val_losses, val_f1s, val_weighted_f1s = [], [], [], []
start_time = time.time()
for epoch in range(1, num_epochs + 1):
@@ -97,21 +97,23 @@ def train_model(
val_loss = val_running_loss / n_val
val_losses.append(val_loss)
- val_f1 = f1_score(all_labels, all_preds, average='weighted', zero_division=0)
- val_f1s.append(val_f1)
+ val_f1_macro = f1_score(all_labels, all_preds, average='macro', zero_division=0)
+ val_f1_weighted = f1_score(all_labels, all_preds, average='weighted', zero_division=0)
+ val_f1s.append(val_f1_macro)
+ val_weighted_f1s.append(val_f1_weighted)
current_lr = scheduler.get_last_lr()[0] if hasattr(scheduler, 'get_last_lr') else optimizer.param_groups[0]['lr']
print(
f"[{model_name}] Epoch {epoch:3d}/{num_epochs} | "
f"train loss: {train_loss:.4f} | val loss: {val_loss:.4f} | "
- f"val F1: {val_f1:.4f} | lr: {current_lr:.2e}"
+ f"macro F1: {val_f1_macro:.4f} | weighted F1: {val_f1_weighted:.4f} | lr: {current_lr:.2e}"
)
scheduler.step()
- # ── Early stopping ────────────────────────────────────────────────────
- if val_f1 > best_val_f1:
- best_val_f1 = val_f1
+ # ── Early stopping (monitored on macro F1) ───────────────────────────
+ if val_f1_macro > best_val_f1:
+ best_val_f1 = val_f1_macro
best_state = copy.deepcopy(model.state_dict())
patience_counter = 0
else:
@@ -123,14 +125,15 @@ def train_model(
train_time = time.time() - start_time
epochs_trained = len(train_losses)
- print(f" Training complete: {epochs_trained} epochs, {train_time:.1f}s, best val F1: {best_val_f1:.4f}")
+ print(f" Training complete: {epochs_trained} epochs, {train_time:.1f}s, best val macro F1: {best_val_f1:.4f}")
return {
'best_model_state': best_state,
'train_losses': train_losses,
'val_losses': val_losses,
- 'val_f1s': val_f1s,
+ 'val_f1s': val_f1s, # macro F1 per epoch (used for early stopping)
+ 'val_weighted_f1s': val_weighted_f1s,
'train_time': train_time,
'epochs_trained': epochs_trained,
- 'best_val_f1': best_val_f1,
+ 'best_val_f1': best_val_f1, # best macro F1
}
diff --git 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()
