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 --- dataset.py | 70 ++++++++++++++++++++++++++++++++++++++++++++----------------- evaluate.py | 34 +++++++++++++++++++----------- main.py | 60 +++++++++++++++++++++++++++++++++++++--------------- train.py | 29 +++++++++++++------------ utils.py | 30 ++++++++++++++++++++++++-- 5 files changed, 160 insertions(+), 63 deletions(-) diff --git a/dataset.py b/dataset.py index 696a312..0068a97 100644 --- 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 index 7a39db4..7f996dc 100644 --- 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 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} | " diff --git a/train.py b/train.py index 545b299..4a04a30 100644 --- 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 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() -- cgit v1.2.3