hpo.py

from MLP import MLP
from CNN import ResCNN
import random
import torch.nn as nn
import torch
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, random_split, Subset
from torch.utils.tensorboard import SummaryWriter
from sklearn.metrics import precision_score, recall_score
from tqdm import tqdm
import matplotlib.pyplot as plt
import time
from datetime import datetime 
import os
import math
import numpy as np
from scipy import interpolate

class HPO:
    def __init__(self, resource_type, domains, metric_to_monitor = "loss", monitor_mode = min, dataset_name = 'mnist', time_unit = 1, eval_frequency = 5, n_runs = 1):
        self.resource_type = resource_type
        self.domains = domains
        self.metric_to_monitor = metric_to_monitor
        self.monitor_mode = monitor_mode
        self.time_unit = time_unit
        self.eval_frequency = eval_frequency  # Evaluate every N epochs
        self.n_runs = n_runs  # Number of runs to average over
         
        self.inf = float('inf') if self.monitor_mode == min else float('-inf')

        self.history = self._init_hist()        
        self.total_resources_used = 0
        self.runs_data = []  # Store data from all runs

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dataset_name = dataset_name
        self.input_size = 28 * 28
        self.input_shape = (1, 28, 28)  # Default for MNIST, will be updated for CIFAR10
        self.output_size = 10
        self.train, self.val = self._load_data()
        self.eval_batch_size = 128  #this should probably stay fixed for evaluation
        self.val_loader = DataLoader(self.val, batch_size = self.eval_batch_size)
        
        self.run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_dir = f"runs/{self.run_id}" #self.__class__.__name__
        os.makedirs(self.log_dir, exist_ok = True)
        self.writer = SummaryWriter(self.log_dir)

        # Track current run state
        self.current_run_history = None
        self.current_best_metric = None
        self.current_best_config = None
        self.current_best_model = None
        self.run_epochs = 0
        self.run_time = 0

    def optimize(self):
        self.history = self._init_hist()
        self.runs_data = []
        self.total_resources_used = 0
        
        for run in range(self.n_runs):
            print(f"[INFO] - Starting run {run + 1}/{self.n_runs}")
            
            self.current_run_history = {
                "metrics": [],           # List of best metrics seen so far
                "timestamps": [],        # List of evaluation timestamps
                "global_epochs": [],     # List of global epochs at evaluation
                "new_config_flags": [],  # List of flags indicating new config evaluation
                "total_resources": 0     # Total resources used in this run
            }
            self.current_best_metric = self.inf
            self.current_best_config = None
            self.current_best_model = None
            self.run_epochs = 0
            self.run_time = 0
            
            self._run_optimization()
            
            self.runs_data.append({
                'history': self.current_run_history.copy(),
                'best_config': self.current_best_config,
                'best_model': self.current_best_model,
                'best_metric': self.current_best_metric
            })
        
        self._aggregate_runs()
        return self.history  
          
    def _run_optimization(self):
        pass
    
    def _register_evaluation(self, metric_value, timestamp, global_epoch, new_config_flag = False, config = None, model = None):
        if self.monitor_mode == min:
            is_better = metric_value < self.current_best_metric
        else:
            is_better = metric_value > self.current_best_metric
            
        if is_better:
            self.current_best_metric = metric_value
            if config is not None:
                self.current_best_config = config.copy()
            if model is not None:
                self.current_best_model = model
        
        self.current_run_history["metrics"].append(self.current_best_metric)
        self.current_run_history["timestamps"].append(timestamp)
        self.current_run_history["global_epochs"].append(global_epoch)
        self.current_run_history["new_config_flags"].append(new_config_flag)
    
    def _aggregate_runs(self):
        if not self.runs_data:
            return
        
        max_timestamp = max(run['history']['timestamps'][-1] for run in self.runs_data if run['history']['timestamps'])
        max_epochs = max(run['history']['global_epochs'][-1] for run in self.runs_data if run['history']['global_epochs'])
        
        max_history_length = max(len(run['history']['metrics']) for run in self.runs_data)

        time_grid = np.linspace(0, max_timestamp, max_history_length)
        epoch_grid = np.linspace(0, max_epochs, max_history_length)
        
        self.history = {
            "time_grid": time_grid,
            "epoch_grid": epoch_grid,
            "metrics_avg": [],
            "metrics_min": [],
            "metrics_max": [],
            "new_config_points_avg": [],  # Average timestamps of new config evaluations
            "new_config_epochs_avg": [],  # Average epochs of new config evaluations
            "total_resources_avg": 0,
            "best_config": None,
            "best_model": None,
            "best_metric": self.inf
        }
        
        interpolated_metrics_time = []
        interpolated_metrics_epoch = []
        
        for run in self.runs_data:
            if not run['history']['timestamps']:
                continue
            
            # Time-based interpolation
            run_interpolated_time = []
            for grid_time in time_grid:
                interpolated_value = self._interpolate_at_point(
                    grid_time, 
                    run['history']['timestamps'], 
                    run['history']['metrics']
                )
                run_interpolated_time.append(interpolated_value)
            
            # Epoch-based interpolation
            run_interpolated_epoch = []
            for grid_epoch in epoch_grid:
                interpolated_value = self._interpolate_at_point(
                    grid_epoch, 
                    run['history']['global_epochs'], 
                    run['history']['metrics']
                )
                run_interpolated_epoch.append(interpolated_value)
            
            interpolated_metrics_time.append(run_interpolated_time)
            interpolated_metrics_epoch.append(run_interpolated_epoch)
            
            # Track overall best
            if self.monitor_mode == min:
                is_better = run['best_metric'] < self.history["best_metric"]
            else:
                is_better = run['best_metric'] > self.history["best_metric"]
                
            if is_better:
                self.history["best_metric"] = run['best_metric']
                self.history["best_config"] = run['best_config']
                self.history["best_model"] = run['best_model']
        
        if interpolated_metrics_time:
            interpolated_metrics_time = np.array(interpolated_metrics_time)
            self.history["metrics_avg"] = np.nanmean(interpolated_metrics_time, axis=0).tolist()
            self.history["metrics_min"] = np.nanmin(interpolated_metrics_time, axis=0).tolist()
            self.history["metrics_max"] = np.nanmax(interpolated_metrics_time, axis=0).tolist()
        
        if interpolated_metrics_epoch:
            interpolated_metrics_epoch = np.array(interpolated_metrics_epoch)
            self.history["metrics_avg_epoch"] = np.nanmean(interpolated_metrics_epoch, axis=0).tolist()
            self.history["metrics_min_epoch"] = np.nanmin(interpolated_metrics_epoch, axis=0).tolist()
            self.history["metrics_max_epoch"] = np.nanmax(interpolated_metrics_epoch, axis=0).tolist()
        
        new_config_timestamps = []
        new_config_epochs = []
        
        for run in self.runs_data:
            run_new_config_times = []
            run_new_config_epochs = []
            
            for i, is_new_config in enumerate(run['history']['new_config_flags']):
                if is_new_config:
                    run_new_config_times.append(run['history']['timestamps'][i])
                    run_new_config_epochs.append(run['history']['global_epochs'][i])
            
            new_config_timestamps.append(run_new_config_times)
            new_config_epochs.append(run_new_config_epochs)
        
        # Average new config points (handling different lengths)
        max_new_configs = max(len(run_configs) for run_configs in new_config_timestamps) if new_config_timestamps else 0
        
        for i in range(max_new_configs):
            time_points = [run_configs[i] for run_configs in new_config_timestamps if i < len(run_configs)]
            epoch_points = [run_configs[i] for run_configs in new_config_epochs if i < len(run_configs)]
            
            if time_points:
                self.history["new_config_points_avg"].append(np.mean(time_points))
            if epoch_points:
                self.history["new_config_epochs_avg"].append(np.mean(epoch_points))
        
        # Average total resources
        self.history["total_resources_avg"] = np.mean([run['history']['total_resources'] for run in self.runs_data])
        self.total_resources_used = self.history["total_resources_avg"]
    
    def _interpolate_at_point(self, target_x, x_values, y_values):
        if len(x_values) == 0:
            return np.nan
        
        # If only one point, return constant value
        if len(x_values) == 1:
            return y_values[0]
        
        # If target is before first point, return first value
        if target_x <= x_values[0]:
            return y_values[0]
        
        # If target is after last point, return last value
        if target_x >= x_values[-1]:
            return y_values[-1]
        
        # Find the interval that contains target_x
        for i in range(len(x_values) - 1):
            if x_values[i] <= target_x <= x_values[i + 1]:
                # Linear interpolation
                x1, x2 = x_values[i], x_values[i + 1]
                y1, y2 = y_values[i], y_values[i + 1]
                
                # Handle case where x1 == x2 (shouldn't happen but just in case)
                if x1 == x2:
                    return y1
                
                # Linear interpolation formula
                return y1 + (y2 - y1) * (target_x - x1) / (x2 - x1)
        
        # Should not reach here, but return last value as fallback
        return y_values[-1]

    def _init_hist(self):
        return {
            "time_grid": [],
            "epoch_grid": [],
            "metrics_avg": [],
            "metrics_min": [],
            "metrics_max": [],
            "new_config_points_avg": [],
            "new_config_epochs_avg": [],
            "total_resources_avg": 0,
            "best_config": None,
            "best_model": None,
            "best_metric": self.inf
        }

    def _load_data(self):
        match(self.dataset_name):
            case 'mnist':
                transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
                dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
                self.input_size = 28 * 28
                self.input_shape = (1, 28, 28)
                self.output_size = 10
                
            case 'cifar10':
                transform = transforms.Compose([
                    transforms.RandomHorizontalFlip(p = 0.5),
                    transforms.RandomRotation(10),
                    transforms.RandomCrop(32, padding = 4),
                    transforms.ToTensor(),
                    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.2435, 0.2616))
                ])
                                
                dataset = torchvision.datasets.CIFAR10(root = './data', train = True, download = True, transform = transform)
                
                self.input_size = 32 * 32 * 3 # 32x32 images with 3 color channels
                self.input_shape = (3, 32, 32)
                self.output_size = 10

            case _ :
                raise ValueError(f"Unknown dataset {self.dataset}")

        size = len(dataset)
        train_size = int(0.8 * size) #80% training, 20% validation
        val_size = size - train_size
        train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
        return train_dataset, val_dataset
        
    def _get_random_config(self, k = 0): 
        if k != 0:
            return self._get_random_configs(k)
        
        config = {}
        for parameter, domain in self.domains.items():
            if isinstance(domain, tuple):
                if len(domain) == 3:
                    min_value, max_value, domain_type = domain
                    if domain_type == "continuous":
                        config[parameter] = random.uniform(min_value, max_value)
                    elif domain_type == "integer":
                        config[parameter] = random.randint(min_value, max_value)
                elif len(domain) == 2: #categorical
                    options, domain_type = domain
                    config[parameter] = random.choice(options)
        return config

    def _get_random_configs(self, k):
        return [self._get_random_config() for _ in range(k)]

    def _training_pipeline(self, model, config, resources, samples = None, save_log = True):          
        train_dataset = self.train
        if samples is not None: train_dataset = Subset(self.train, list(range(samples))) 
        else: samples = len(train_dataset)

        batch_size = config.get("batch_size", min(128, samples) )
        train_loader = DataLoader(train_dataset, batch_size = batch_size, shuffle = True)

        lr = config.get("learning_rate", 0.001)
        optimizer = config.get("optimizer", optim.Adam)(model.parameters(), lr = lr)
        criterion = config.get("loss", nn.CrossEntropyLoss()) #this should probably be fixed
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor = 0.5, patience = 3, threshold = 1e-4)
        
        model.to(self.device)
        model.train()
        
        hist = {
            "accuracy": [],
            "precision": [],
            "recall": [],
            "loss": [],
            "test_error": []
        }


        epoch = 0
        previous_time = time.perf_counter()
        elapsed = 0

        is_new_config = False

        while True:
            model.train()

            epoch_preds, epoch_labels = [], []
            epoch_loss = 0.0

            if epoch == 0:
                is_new_config = True
            else:
                is_new_config = False

            for X, y in train_loader:
                X, y = X.to(self.device), y.to(self.device)
                optimizer.zero_grad()
                
                out = model(X)
                loss = criterion(out, y)
                loss.backward()
                optimizer.step()

                preds = out.argmax(dim = 1).cpu().tolist()
                epoch_preds += preds
                epoch_labels += y.cpu().tolist()

                epoch_loss += loss.item() * X.size(0)
            

            n = len(epoch_labels)
            acc = sum(p == t for p, t in zip(epoch_preds, epoch_labels)) / n
            prec = precision_score(epoch_labels, epoch_preds, average = "macro", zero_division = 0)
            rec  = recall_score(epoch_labels, epoch_preds, average = "macro", zero_division = 0)
            loss_epoch = epoch_loss / n
            test_err = 1.0 - acc

            hist["accuracy"].append(acc)
            hist["precision"].append(prec)
            hist["recall"].append(rec)
            hist["loss"].append(loss_epoch)
            hist["test_error"].append(test_err)
            
            scheduler.step(loss_epoch)
        
            # Evaluate on full validation set periodically and register in history
            if save_log and (is_new_config or epoch % self.eval_frequency == 0 or 
                           (self.resource_type == "epochs" and epoch >= resources) or #case is last time we evaluate the config
                           (self.resource_type == "time" and elapsed > resources * self.time_unit - last_time)):
                
                val_metrics = self._evaluate_model(model)
                val_metric_value = val_metrics[self.metric_to_monitor]
                
                self._register_evaluation(
                    val_metric_value, 
                    self.run_time, 
                    self.run_epochs, 
                    new_config_flag = is_new_config,
                    config = config,
                    model = model
                )
                
                # Log to tensorboard
                self.writer.add_scalar("best", self.current_best_metric, self.run_epochs)

            epoch += 1

            current_time = time.perf_counter()
            last_time = current_time - previous_time
            previous_time = current_time

            elapsed += last_time

            if save_log:
                self.run_epochs += 1
                self.run_time += last_time

            if (self.resource_type == "epochs" and epoch >= resources) or (self.resource_type == "time" and elapsed > resources*self.time_unit - last_time):
                break
 
        if save_log:
            match self.resource_type:
                case "epochs": used = epoch
                case "time": used = elapsed

            self.current_run_history["total_resources"] += used
            self.total_resources_used += used
            print(f"[INFO] - used resources = {used}, total_resources = {self.total_resources_used}")
        
        return hist
    
    def _evaluate_model(self, model):
        model.eval()
        all_preds, all_labels = [], []
        total_loss = 0.0
        criterion = nn.CrossEntropyLoss()

        with torch.no_grad():
            for X, y in self.val_loader:
                X, y = X.to(self.device), y.to(self.device)
                out = model(X)
                loss = criterion(out, y)
                total_loss += loss.item() * X.size(0)
                preds = out.argmax(dim = 1).cpu().tolist()
                all_preds += preds
                all_labels += y.cpu().tolist()
       
        n = len(all_labels)
        accuracy  = sum(int(p==t) for p,t in zip(all_preds, all_labels)) / n
        precision = precision_score(all_labels, all_preds, average = 'macro', zero_division = 0)
        recall    = recall_score(all_labels, all_preds, average = 'macro', zero_division = 0)
        loss      = total_loss / n
        test_error = 1. - accuracy
        
        model.train() 

        return {
            "accuracy": accuracy,
            "precision": precision,
            "recall": recall,
            "loss": loss,
            "test_error": test_error
        }

    def build_model(self, config):
        #print(config)

        if self.dataset_name == "mnist":
            num_layers = config.get("num_layers", 1)
            hidden_units = config.get("hidden_units", 64)
            dropout_rate = config.get("dropout_rate", 0.0)
            act_funcs = config.get("activation_function", nn.ReLU)

            layer_sizes = [self.input_size] + [hidden_units] * num_layers + [self.output_size]
            mlp = MLP(layer_sizes=layer_sizes, dropouts=dropout_rate, act_funcs=act_funcs)
            return mlp

        depth = config.get("depth", 12)
        channel_progression = config.get("channel_progression", [64, 128, 256, 512])
        use_batch_norm = config.get("use_batch_norm", True)
        kernel = config.get("kernel", 3)
        stride = 1
        padding = kernel // 2 
        pool_kernel = config.get("pool_kernel", 2)
        pool_stride = 2
        pool_padding = 0 if pool_kernel == 2 else 1
        act_funcs = config.get("activation_function", nn.ReLU)
        dropout_rate = config.get("dropout_rate", 0.0)
        mlp_hidden_dim = config.get("mlp_hidden_dim", 512)
        mlp_dropout = config.get("mlp_dropout", 0.5)

        input_h, input_w = self.input_shape[1:]
        current_h, current_w = input_h, input_w

        architecture = []

        max_pool_layers = int(math.log2(min(current_h, current_w)) - 2)
        blocks = max(1, max_pool_layers)
        base_num_conv = depth // blocks
        extras = depth % blocks

        layer_id = 0
        for block in range(blocks):
            num_conv = base_num_conv + (1 if block < extras else 0)
            out_chan = channel_progression[block] if block < len(channel_progression) else channel_progression[-1]

            for conv in range(num_conv):
                layer_config = {
                    "out_channels": out_chan,
                    "kernel": kernel,
                    "stride": stride,
                    "padding": padding,
                    "use_batch_norm": use_batch_norm,
                    "act_func": act_funcs,
                }

                if dropout_rate > 0.0:
                    layer_config["dropout"] = dropout_rate

                # Pooling on last conv in block
                if conv == num_conv - 1:
                    est_h = (current_h + 2 * pool_padding - pool_kernel) // pool_stride + 1
                    est_w = (current_w + 2 * pool_padding - pool_kernel) // pool_stride + 1

                    if est_h >= 1 and est_w >= 1:
                        layer_config["pool"] = nn.MaxPool2d
                        layer_config["pool_kernel"] = pool_kernel
                        layer_config["pool_stride"] = pool_stride
                        layer_config["pool_padding"] = pool_padding
                        current_h, current_w = est_h, est_w
                    else:
                        print(f"[WARNING] Skipped pooling at layer {layer_id} to avoid collapse ({current_h}x{current_w})")

                architecture.append(layer_config)
                layer_id += 1

        norm_layer = nn.BatchNorm2d if use_batch_norm else nn.Identity

        rescnn = ResCNN(self.input_shape, architecture, norm_layer=norm_layer)
        rescnn_output_size = rescnn.get_output_size()
        #print(f"[INFO] - ResCNN output size: {rescnn_output_size}")

        mlp_tail = MLP(
            layer_sizes=[rescnn_output_size, mlp_hidden_dim, self.output_size],
            act_funcs=[act_funcs, None],
            dropouts=[mlp_dropout, 0.0],
            use_bias=True,
        )
        rescnn.attach_mlp(mlp_tail)
        return rescnn
    
    def plot(self, custom_y_scale = True):
        # Plot based on time or epochs
        plt.figure(figsize=(10, 6))
            
        #if self.resource_type == 'time':
        x_values = self.history["time_grid"]
        metrics_avg = self.history["metrics_avg"]
        metrics_min = self.history["metrics_min"]
        metrics_max = self.history["metrics_max"]
        new_config_points = self.history["new_config_points_avg"]
        xlabel = "Time (seconds)"
        #else:
        #    x_values = self.history["epoch_grid"]
        #    metrics_avg = self.history.get("metrics_avg_epoch", [])
        #    metrics_min = self.history.get("metrics_min_epoch", [])
        #    metrics_max = self.history.get("metrics_max_epoch", [])
        #    new_config_points = self.history["new_config_epochs_avg"]
        #    xlabel = "Global Epochs"
            
        if metrics_avg and len(x_values) == len(metrics_avg):
            # Transform data if custom y-scale is requested
            if custom_y_scale:
                metrics_avg_plot = [self._transform_y_value(y) for y in metrics_avg]
                metrics_min_plot = [self._transform_y_value(y) for y in metrics_min] if metrics_min else []
                metrics_max_plot = [self._transform_y_value(y) for y in metrics_max] if metrics_max else []
            else:
                metrics_avg_plot = metrics_avg
                metrics_min_plot = metrics_min
                metrics_max_plot = metrics_max
            
            # Plot average with min/max area
            plt.plot(x_values, metrics_avg_plot, label="Average", color='blue', linewidth=2)
            if metrics_min_plot and metrics_max_plot:
                plt.fill_between(x_values, metrics_min_plot, metrics_max_plot, alpha=0.3, color='blue', label="Min-Max Range")
           
            if new_config_points:
                # Find corresponding y-values for new config points
                new_config_x = []
                new_config_y = []
                for point in new_config_points:
                    # Find closest x-value and corresponding y-value
                    closest_idx = min(range(len(x_values)), key=lambda i: abs(x_values[i] - point))
                    x_point = x_values[closest_idx]
                    y_point = metrics_avg_plot[closest_idx]
                    new_config_x.append(x_point)
                    new_config_y.append(y_point)

                    plt.plot([0, x_point], [y_point, y_point], '--', alpha=0.7,color = 'lightgray', linewidth = 1)  # vertical line to x-axis
                    plt.plot([x_point, x_point], [self._transform_y_value(0), y_point], '--', alpha=0.7,color = 'lightgray', linewidth = 1)
                
                plt.scatter(new_config_x, new_config_y, color="blue", s = 50, zorder=5, 
                           label="New Config Evaluations", marker='o')

        
        # Apply custom y-scale if requested
        if custom_y_scale:
            self._apply_custom_y_scale(plt.gca())
        else:
            plt.grid(True, color='lightgray', linestyle='-', alpha=0.7, linewidth=0.5)
        
        plt.xlabel(xlabel)
        plt.ylabel(self.metric_to_monitor)
        plt.title(f"{self.__class__.__name__} Results")
        plt.legend()
        plt.tight_layout()
        plt.savefig(f"{self.resource_type}_{self.__class__.__name__}_{self.run_id}.png", dpi=300, bbox_inches='tight')
        plt.close()

    def _transform_y_value(self, y):
        """Transform a single y-value for custom scaling"""
        if y >= 0.1:
            return 0.5 + 0.5 * (y - 0.1) / (1.0 - 0.1)
        else:
            return 0.5 * (y - 0.01) / (0.1 - 0.01)

    def _apply_custom_y_scale(self, ax):
        """Apply custom y-axis scaling with proper tick labels and gridlines"""
        # Define the tick positions and labels
        y_ticks = [1.0, 0.5, 0.1, 0.05, 0.01]
        y_tick_positions = [self._transform_y_value(tick) for tick in y_ticks]
        y_tick_labels = ['10^0', '0.5', '10^-1', '0.05', '10^-2']
        
        ax.set_yticks(y_tick_positions)
        ax.set_yticklabels(y_tick_labels)
        ax.set_ylim(0, 1)