6_hyperparameter_tuning.py

import pandas as pd
import numpy as np
from sklearn.metrics import mean_absolute_error
from torch.utils.data import Dataset, DataLoader
import torch
import timm
import torch.nn as nn
from PIL import Image
from torchvision import transforms
from tqdm import tqdm
import os
import time
from datetime import datetime
import matplotlib.pyplot as plt
from torchvision.utils import make_grid
import math
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay, f1_score, balanced_accuracy_score, mean_absolute_error, mean_squared_error
import shutil
from collections import OrderedDict
import gc
from sklearn.model_selection import StratifiedKFold
from scipy import stats
import json
import pickle
import itertools
from torch.optim.lr_scheduler import StepLR, CosineAnnealingLR, ReduceLROnPlateau
 
def cleanup_after_fold():
    # Delete key objects
    del model
    del optimizer
    if 'scheduler' in locals() or 'scheduler' in globals():
        del scheduler
   
    # Close dataloaders
    for loader in [train_loader, val_loader]:
        loader._iterator = None
        if hasattr(loader, '_workers'):
            for w in loader._workers:
                if w.is_alive():
                    w.terminate()
 
# Define root_mean_squared_error function since it's not in sklearn
def root_mean_squared_error(y_true, y_pred):
    return np.sqrt(mean_squared_error(y_true, y_pred))
 
# Enable more detailed logging
def log_separator(message="", char="=", length=80):
    """Print a separator line with optional message"""
    if message:
        message = f" {message} "
    padding = (length - len(message)) // 2
    print(f"\n{char * padding}{message}{char * padding}")
 
# Checkpoint functions for hyperparameter tuning
def save_hp_checkpoint(output_path, model_name, completed_configs, current_config_idx, all_results):
    """Save checkpoint during hyperparameter tuning"""
    checkpoint_data = {
        'model_name': model_name,
        'completed_configs': completed_configs,
        'current_config_idx': current_config_idx,
        'all_results': all_results,
        'timestamp': datetime.now().isoformat()
    }
   
    checkpoint_path = os.path.join(output_path, 'hp_checkpoint.json')
    with open(checkpoint_path, 'w') as f:
        json.dump(checkpoint_data, f, indent=2)
   
    print(f"✓ Hyperparameter checkpoint saved: completed {len(completed_configs)} configurations")
 
def load_hp_checkpoint(output_path):
    """Load hyperparameter checkpoint if exists"""
    checkpoint_path = os.path.join(output_path, 'hp_checkpoint.json')
    if os.path.exists(checkpoint_path):
        with open(checkpoint_path, 'r') as f:
            checkpoint_data = json.load(f)
        print(f"✓ Hyperparameter checkpoint found: resuming from configuration {checkpoint_data['current_config_idx']}")
        return checkpoint_data
    return None
 
class PDFFDataset(Dataset):
    def __init__(self, df, transform=None, augment_training=False, model_name='resnet18', resize_pixel=56):
        self.df = df
        self.augment_training = augment_training
       
        # Get the input size for the model
        self.model_name = model_name
        self.input_size = self._get_model_input_size(model_name)
 
        if resize_pixel is None:
            self.resize_pixel = self.input_size
        else:
            self.resize_pixel = resize_pixel
               
        # Base transform for all images (crop, resize, convert to tensor, and normalize)
        self.base_transform = transforms.Compose([
            transforms.Lambda(lambda img: self._crop_borders(img, border_size=50)),  # Crop 50 pixels from all borders
            transforms.Resize(self.resize_pixel),
            transforms.Resize(self.input_size),  # Resize to model's input size
            transforms.ToTensor()
        ])
       
        # Custom transform if provided
        self.custom_transform = transform
 
        # Specific augmentations
        self.h_flip = transforms.RandomHorizontalFlip(p=1.0)  # Always flip horizontally
        self.rotate = transforms.Lambda(lambda img: self._rotate_from_top_center(img, angle=20))  # Custom rotation
 
        print(f"Dataset created with {len(self.df)} original samples")
        if augment_training:
            total_samples = len(self)
            original_samples = len(self.df)
            print(f"Augmentation enabled - Total samples: {total_samples}")
            print(f"   → Original: {original_samples}")
            print(f"   → Horizontal flips: {original_samples}")
            print(f"   → Rotated: {original_samples}")
        else:
            print(f"No augmentation - Total samples: {len(self)}")
 
    def _get_model_input_size(self, model_name):
        """Returns the correct input size for the specified model."""
        try:
            model = timm.create_model(model_name, pretrained=True)
            if hasattr(model, 'default_cfg') and 'input_size' in model.default_cfg:
                input_size = model.default_cfg['input_size'][1:]  # Returns (3, H, W), we need (H, W)
            else:
                input_size = (224, 224)
        except Exception as e:
            print(f"Error creating model {model_name}: {e}")
            input_size = (224, 224)
        return input_size
 
    def _crop_borders(self, img, border_size=50):
        """Crop specified number of pixels from all borders of the image"""
        width, height = img.size
        new_width = width - (2 * border_size)
        new_height = height - (2 * border_size)
        if new_width <= 0 or new_height <= 0:
            print(f"Warning: Image too small to crop {border_size} pixels from borders. Original size: {width}x{height}")
            return img
        return img.crop((border_size, border_size, width - border_size, height - border_size))
       
    def _rotate_from_top_center(self, img, angle=15):
        """Rotate the image around the top-center point"""
        import math
        width, height = img.size
        angle_rad = math.radians(angle)
 
        cx, cy = width / 2, 0  # top center
 
        cos_theta = math.cos(angle_rad)
        sin_theta = math.sin(angle_rad)
 
        # Affine transform matrix values for PIL (a, b, c, d, e, f)
        a = cos_theta
        b = sin_theta
        c = cx - cos_theta * cx - sin_theta * cy
        d = -sin_theta
        e = cos_theta
        f = cy - (-sin_theta) * cx - cos_theta * cy
 
        return img.transform(
            img.size,
            Image.AFFINE,
            (a, b, c, d, e, f),
            resample=Image.BICUBIC,
            fillcolor=(0, 0, 0)
        )
 
    def __len__(self):
        if self.augment_training:
            return len(self.df) * 3
        return len(self.df)
 
    def __getitem__(self, idx):
        if self.augment_training:
            real_idx = idx // 3
            aug_type = idx % 3
        else:
            real_idx = idx
            aug_type = 0
           
        row = self.df.iloc[real_idx]
        img_path = row["img_path"]
       
        try:
            img = Image.open(img_path).convert("RGB")
        except Exception as e:
            print(f"Error loading image {img_path}: {e}")
            raise
           
        if self.custom_transform:
            img = self.custom_transform(img)
        else:
            if self.augment_training:
                if aug_type == 1:  # Horizontal flip only
                    img = self.h_flip(img)
                elif aug_type == 2:  # Rotation only
                    img = self.rotate(img)
 
            img = self.base_transform(img)
       
        label = torch.tensor(row["PDFF Fat fraction"], dtype=torch.float32)
        return img, label
 
def train_one_epoch(model, dataloader, optimizer, loss_fn, scheduler, epoch_num):
    model.train()
    running_loss = 0.0
    total_samples = 0
    batch_count = 0
    start_time = time.time()
    progress_bar = tqdm(dataloader, desc=f"Epoch {epoch_num} training")
   
    for x, y in progress_bar:
        batch_size = x.size(0)
        total_samples += batch_size
        batch_count += 1
       
        x, y = x.cuda(), y.cuda().unsqueeze(1)
        optimizer.zero_grad()
        preds = model(x)
        loss = loss_fn(preds, y)
        loss.backward()
        optimizer.step()
       
        # Update running statistics
        running_loss += loss.item() * batch_size
        progress_bar.set_postfix({"loss": f"{running_loss/total_samples:.4f}"})
 
    # Step scheduler if it's not ReduceLROnPlateau
    if scheduler is not None and not isinstance(scheduler, ReduceLROnPlateau):
        scheduler.step()
 
    epoch_loss = running_loss / total_samples
    epoch_time = time.time() - start_time
   
    print(f"Epoch {epoch_num} completed in {epoch_time:.2f}s")
    print(f"  • Batches: {batch_count}, Total samples: {total_samples}")
    print(f"  • Average loss: {epoch_loss:.4f}")
   
    return epoch_loss
 
def evaluate(model, dataloader, phase="val"):
    model.eval()
    preds, targets = [], []
    running_loss = 0.0
    loss_fn = nn.L1Loss()
    total_samples = 0
   
    start_time = time.time()
    progress_bar = tqdm(dataloader, desc=f"{phase.capitalize()} evaluation")
   
    with torch.no_grad():
        for x, y in progress_bar:
            batch_size = x.size(0)
            total_samples += batch_size
           
            x = x.cuda()
            y_tensor = y.cuda().unsqueeze(1)
           
            # Forward pass
            batch_preds = model(x)
           
            # Calculate loss
            loss = loss_fn(batch_preds, y_tensor)
            running_loss += loss.item() * batch_size
           
            # Store predictions and targets
            preds.append(batch_preds.cpu().numpy())
            targets.append(y.numpy().reshape(-1, 1))
           
            # Update progress bar
            progress_bar.set_postfix({"loss": f"{running_loss/total_samples:.4f}"})
   
    # Concatenate results
    preds_array = np.concatenate(preds)
    targets_array = np.concatenate(targets)
   
    # Calculate metrics
    mae = mean_absolute_error(targets_array, preds_array)
    avg_loss = running_loss / total_samples
    eval_time = time.time() - start_time
   
    print(f"{phase.capitalize()} completed in {eval_time:.2f}s")
    print(f"  • Total samples: {total_samples}")
    print(f"  • L1 Loss: {avg_loss:.4f}")
    print(f"  • MAE: {mae:.4f}")
   
    return preds_array, targets_array
 
def get_scheduler(optimizer, scheduler_type, **kwargs):
    """Create scheduler based on type"""
    if scheduler_type == 'step':
        return StepLR(optimizer, step_size=kwargs.get('step_size', 30), gamma=kwargs.get('gamma', 0.1))
    elif scheduler_type == 'cosine':
        return CosineAnnealingLR(optimizer, T_max=kwargs.get('t_max', 50))
    elif scheduler_type == 'plateau':
        return ReduceLROnPlateau(optimizer, mode='min', factor=kwargs.get('factor', 0.5),
                               patience=kwargs.get('patience', 10), verbose=True)
    else:
        return None
 
def run_single_fold_validation(model_name, config, cv_df, fold_idx=0, num_epochs=30):
    """Run a single fold validation for hyperparameter tuning"""
   
    # Split data for this fold
    train_data = cv_df[cv_df["fold"] != fold_idx].reset_index(drop=True)
    val_data = cv_df[cv_df["fold"] == fold_idx].reset_index(drop=True)
   
    # Create datasets
    train_dataset = PDFFDataset(train_data, augment_training=True, model_name=model_name)
    val_dataset = PDFFDataset(val_data, augment_training=False, model_name=model_name)
   
    # Create dataloaders
    train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=0)
    val_loader = DataLoader(val_dataset, batch_size=config['batch_size'], num_workers=0)
   
    # Model setup
    model = timm.create_model(model_name, pretrained=True, num_classes=1).cuda()
   
    # Optimizer setup
    if config['optimizer'] == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=config['learning_rate'],
                                   weight_decay=config['weight_decay'])
    elif config['optimizer'] == 'adamw':
        optimizer = torch.optim.AdamW(model.parameters(), lr=config['learning_rate'],
                                    weight_decay=config['weight_decay'])
    elif config['optimizer'] == 'sgd':
        optimizer = torch.optim.SGD(model.parameters(), lr=config['learning_rate'],
                                  weight_decay=config['weight_decay'], momentum=0.9)
   
    # Scheduler setup
    scheduler = get_scheduler(optimizer, config['scheduler'],
                            step_size=config.get('step_size', 30),
                            gamma=config.get('gamma', 0.1),
                            t_max=num_epochs,
                            factor=config.get('factor', 0.5),
                            patience=config.get('patience', 10))
   
    # Loss function
    if config['loss_function'] == 'l1':
        loss_fn = nn.L1Loss()
    elif config['loss_function'] == 'mse':
        loss_fn = nn.MSELoss()
    elif config['loss_function'] == 'huber':
        loss_fn = nn.HuberLoss(delta=config.get('huber_delta', 1.0))
   
    # Training loop
    best_val_mae = float('inf')
    best_epoch = 0
   
    for epoch in range(num_epochs):
        # Train
        train_loss = train_one_epoch(model, train_loader, optimizer, loss_fn, scheduler, epoch+1)
       
        # Validate
        preds, targets = evaluate(model, val_loader, phase="validation")
        val_mae = mean_absolute_error(targets, preds)
       
        # Step ReduceLROnPlateau scheduler
        if isinstance(scheduler, ReduceLROnPlateau):
            scheduler.step(val_mae)
       
        # Track best model
        if val_mae < best_val_mae:
            best_val_mae = val_mae
            best_epoch = epoch
   
    # Cleanup
    del model, optimizer, scheduler, train_dataset, val_dataset, train_loader, val_loader
    torch.cuda.empty_cache()
    gc.collect()
   
    return best_val_mae, best_epoch
 
def generate_hyperparameter_configs():
    """Generate all hyperparameter combinations to test"""
   
    # Define hyperparameter search space
    param_grid = {
        'learning_rate': [1e-5, 1e-4, 1e-3],
        'weight_decay': [1e-6, 1e-4, 1e-2],
        'batch_size': [8, 16, 32],
        'optimizer': ['adam', 'adamw'],
        'scheduler': ['step', 'cosine', None],
        'loss_function': ['l1', 'mse', 'huber']
    }
   
    # Generate all combinations
    keys = param_grid.keys()
    values = param_grid.values()
    configs = []
   
    for combination in itertools.product(*values):
        config = dict(zip(keys, combination))
       
        # Add scheduler-specific parameters
        if config['scheduler'] == 'step':
            config['step_size'] = 20
            config['gamma'] = 0.5
        elif config['scheduler'] == 'plateau':
            config['factor'] = 0.5
            config['patience'] = 8
        elif config['scheduler'] == 'cosine':
            pass  # T_max will be set to num_epochs
       
        # Add loss-specific parameters
        if config['loss_function'] == 'huber':
            config['huber_delta'] = 1.0
            
        configs.append(config)
   
    return configs
 
def hyperparameter_tuning(model_name, cv_df, output_path, num_epochs=30, top_k=10):
    """Main hyperparameter tuning function"""
   
    log_separator(f"HYPERPARAMETER TUNING FOR {model_name}", "=")
   
    # Create output directory
    os.makedirs(output_path, exist_ok=True)
   
    # Check for existing checkpoint
    checkpoint_data = load_hp_checkpoint(output_path)
   
    if checkpoint_data is not None:
        # Resume from checkpoint
        completed_configs = checkpoint_data['completed_configs']
        current_config_idx = checkpoint_data['current_config_idx']
        all_results = checkpoint_data['all_results']
        all_configs = generate_hyperparameter_configs()
        print(f"Resuming from configuration {current_config_idx}/{len(all_configs)}")
    else:
        # Start fresh
        all_configs = generate_hyperparameter_configs()
        completed_configs = []
        all_results = []
        current_config_idx = 0
        print(f"Starting hyperparameter tuning with {len(all_configs)} configurations")
   
    # Test each configuration
    for i, config in enumerate(all_configs[current_config_idx:], current_config_idx):
        log_separator(f"Configuration {i+1}/{len(all_configs)}", "-")
        print(f"Testing configuration: {config}")
       
        try:
            # Run single fold validation (using fold 0 for speed)
            best_mae, best_epoch = run_single_fold_validation(model_name, config, cv_df,
                                                            fold_idx=0, num_epochs=num_epochs)
           
            result = {
                'config_idx': i,
                'config': config,
                'best_val_mae': best_mae,
                'best_epoch': best_epoch
            }
           
            all_results.append(result)
            completed_configs.append(i)
           
            print(f"✓ Configuration {i+1} completed - Best MAE: {best_mae:.4f} at epoch {best_epoch}")
           
            # Save checkpoint every 10 configurations
            if (i + 1) % 10 == 0:
                save_hp_checkpoint(output_path, model_name, completed_configs, i+1, all_results)
               
                # Save intermediate results
                results_df = pd.DataFrame([r for r in all_results])
                results_df = results_df.sort_values('best_val_mae').reset_index(drop=True)
                results_df.to_csv(f"{output_path}/intermediate_results.csv", index=False)
                print(f"Intermediate results saved. Current best MAE: {results_df.iloc[0]['best_val_mae']:.4f}")
           
        except Exception as e:
            print(f"Error in configuration {i+1}: {e}")
            continue
   
    # Final results processing
    log_separator("HYPERPARAMETER TUNING COMPLETED", "=")
   
    # Convert results to DataFrame and sort by performance
    results_df = pd.DataFrame(all_results)
    results_df = results_df.sort_values('best_val_mae').reset_index(drop=True)
   
    # Save all results
    results_df.to_csv(f"{output_path}/all_hyperparameter_results.csv", index=False)
   
    # Save top k configurations
    top_k_results = results_df.head(top_k)
    top_k_results.to_csv(f"{output_path}/top_{top_k}_configurations.csv", index=False)
   
    # Print top results
    print(f"\nTop {top_k} configurations:")
    for i, row in top_k_results.iterrows():
        print(f"\n{i+1}. MAE: {row['best_val_mae']:.4f}")
        config_str = ", ".join([f"{k}={v}" for k, v in row['config'].items()])
        print(f"   Config: {config_str}")
   
    # Clean up checkpoint file
    checkpoint_path = os.path.join(output_path, 'hp_checkpoint.json')
    if os.path.exists(checkpoint_path):
        os.remove(checkpoint_path)
        print("✓ Hyperparameter checkpoint file removed after completion")
   
    return top_k_results
 
####################################################################################
# MAIN EXECUTION
####################################################################################
 
# Clear CUDA cache
torch.cuda.empty_cache()
gc.collect()
 
# Target models for hyperparameter tuning
target_models = [
    'swinv2_tiny_window8_256.ms_in1k',
    'coatnet_0_rw_224.sw_in1k',
    # 'swin_tiny_patch4_window7_224.ms_in1k',
    # 'volo_d1_224.sail_in1k'
]
 
# Parameters
input_path = '/research/projects/Nico/Liver_US_Classification/'
base_output_path = '/research/projects/Nico/Liver_US_Classification/Hyperparameter_Tuning_Results/'
threshold = 5
num_epochs_hp = 40  # Reduced epochs for hyperparameter search
top_k_configs = 5
 
# Load and prepare data
print("Loading and preparing data...")
df = pd.read_csv(f"{input_path}dataset_v2.csv")
df["pdff_class_bin"] = (df["pdff_class"] >= 1).astype(int)
df = df[df["set_v2"] != "unassigned"].reset_index(drop=True)
 
# Split test set
test_df = df[df["set_v2"] == "test"].reset_index(drop=True)
cv_df = df[df["set_v2"] != "test"].reset_index(drop=True)
 
# Create fold assignments for hyperparameter tuning (using same strategy as original)
log_separator("CREATING FOLD ASSIGNMENTS", "-")
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
cv_df["fold"] = -1
 
for fold, (train_idx, val_idx) in enumerate(skf.split(cv_df, cv_df["pdff_class_bin"])):
    cv_df.loc[val_idx, "fold"] = fold
    print(f"Fold {fold}: {len(train_idx)} train samples, {len(val_idx)} validation samples")
 
# Run hyperparameter tuning for each target model
for model_name in target_models:
    print(f"\n{'='*100}")
    print(f"STARTING HYPERPARAMETER TUNING FOR: {model_name}")
    print(f"{'='*100}")
   
    model_output_path = f"{base_output_path}/{model_name}/"
   
    try:
        # Run hyperparameter tuning
        top_configs = hyperparameter_tuning(
            model_name=model_name,
            cv_df=cv_df,
            output_path=model_output_path,
            num_epochs=num_epochs_hp,
            top_k=top_k_configs
        )
       
        print(f"\n✓ Hyperparameter tuning completed for {model_name}")
        print(f"Best configuration achieves MAE: {top_configs.iloc[0]['best_val_mae']:.4f}")
       
    except Exception as e:
        print(f"Error during hyperparameter tuning for {model_name}: {e}")
        continue
 
print("\n" + "="*100)
print("HYPERPARAMETER TUNING COMPLETED FOR ALL MODELS")
print("="*100)
print("\nNext steps:")
print("1. Review the top configurations in the CSV files")
print("2. Select the best hyperparameters for each model")
print("3. Run full 5-fold cross-validation with the optimal settings")
print("4. Train final models on the full dataset")