octopi 1.4.0__py3-none-any.whl

octopi/pytorch/hyper_search.py

@@ -0,0 +1,244 @@
+from monai.metrics import ConfusionMatrixMetric
+from mlflow.tracking import MlflowClient
+from octopi.pytorch import trainer
+from octopi.models import common
+import torch, mlflow, optuna, gc
+from octopi.utils import io
+
+class BayesianModelSearch:
+
+    def __init__(self, data_generator, model_type="Unet", parent_run_id=None, parent_run_name=None):
+        """
+        Class to handle model creation, training, and optimization.
+
+        Args:
+            data_generator (object): Data generator object containing dataset properties.
+            model_type (str): Type of model to build ("UNet", "AttentionUnet").
+        """
+        self.data_generator = data_generator
+        self.Nclasses = data_generator.Nclasses
+        self.device = None
+        self.model_type = model_type
+        self.model = None
+        self.loss_function = None
+        self.metrics_function = None
+        self.sampling = None
+        self.parent_run_id = parent_run_id
+        
+        # Define results directory path
+        self.results_dir = f'explore_results_{self.model_type}'            
+
+    def my_build_model(self, trial):
+        """Builds and initializes a model based on Optuna-suggested parameters."""
+       
+        # Build the model
+        self.model_builder = common.get_model(self.model_type)
+        self.model_builder.bayesian_search(trial, self.Nclasses)
+        self.model = self.model_builder.model.to(self.device)
+        self.config = self.model_builder.config
+
+        # Define loss function
+        self.loss_function = common.get_loss_function(trial)
+
+        # Define metrics
+        self.metrics_function = ConfusionMatrixMetric(
+            include_background=False,
+            metric_name=["recall", "precision", "f1 score"],
+            reduction="none"
+        )
+
+        # Sample crop size and num_samples
+        self.sampling = {
+            'crop_size': trial.suggest_int("crop_size", 48, 192, step=16),
+            'num_samples': 8
+        }
+        self.config['dim_in'] = self.sampling['crop_size']
+
+    def _define_optimizer(self, trial):
+        # Define optimizer
+        # lr0 = trial.suggest_float("lr", 1e-4, 1e-3, log=True)
+        # wd = trial.suggest_float("weight_decay", 1e-5, 1e-3, log=True)
+        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-3, weight_decay=1e-4)
+
+    def _train_model(self, trial, model_trainer, epochs, val_interval, crop_size, num_samples, best_metric):
+        """Handles model training and error handling."""
+        try:
+            results = model_trainer.train(
+                self.data_generator,
+                model_save_path=None,
+                crop_size=crop_size,
+                max_epochs=epochs,
+                val_interval=val_interval,
+                my_num_samples=num_samples,
+                best_metric=best_metric,
+                use_mlflow=True,
+                verbose=False,
+                trial=trial
+            )
+            return results['best_metric']
+
+        except torch.cuda.OutOfMemoryError:
+            print(f"[Trial Failed] OOM Error for model={model_trainer.model}, crop_size={crop_size}, num_samples={num_samples}")
+            trial.set_user_attr("out_of_memory", True)
+            raise optuna.TrialPruned()
+
+        except Exception as e:
+            print(f"[Trial Failed] Unexpected error: {e}")
+            trial.set_user_attr("error", str(e))
+            raise optuna.TrialPruned()        
+
+    def objective(self, trial, epochs, device, val_interval=15, best_metric="avg_f1"):
+        """Runs the full training process for a given trial."""
+
+        # Set device
+        self.device = device
+        
+        # Set a unique run name for each trial
+        trial_num = f"trial_{trial.number}"
+
+        # Start MLflow run
+        with mlflow.start_run(run_name=trial_num, nested=True):
+            
+            # Build model
+            self.my_build_model(trial)
+
+            # Create trainer
+            self._define_optimizer(trial)
+            model_trainer = trainer.ModelTrainer(self.model, self.device, self.loss_function, self.metrics_function, self.optimizer)
+
+            # Train model and evaluate score
+            score = self._train_model(
+                trial, model_trainer, epochs, val_interval, 
+                self.sampling['crop_size'], self.sampling['num_samples'], 
+                best_metric)               
+
+            # Log parameters and metrics
+            params = {
+                'model': self.model_builder.get_model_parameters(),
+                'optimizer': io.get_optimizer_parameters(model_trainer)
+            }
+            model_trainer.my_log_params(io.flatten_params(params))
+
+            # Explicitly set the parent run ID
+            mlflow.log_param("parent_run_id", self.parent_run_id)
+            mlflow.log_param("parent_run_name", self.parent_run_name)         
+
+            # Save best model
+            self._save_best_model(trial, model_trainer, score)
+
+            # Cleanup
+            self.cleanup(model_trainer)
+
+        return score        
+
+    def _setup_parallel_trial_run(self, trial, parent_run=None, gpu_count=1):
+        """Set up parallel MLflow runs and assign GPU for the trial."""
+        
+        trial_num = f"trial_{trial.number}"
+
+        # Create a child run under the parent MLflow experiment
+        mlflow_client = MlflowClient()
+        self.trial_run = mlflow_client.create_run(
+            experiment_id=mlflow_client.get_run(parent_run.info.run_id).info.experiment_id,
+            tags={"mlflow.parentRunId": parent_run.info.run_id},
+            run_name=trial_num
+        )
+        target_run_id = self.trial_run.info.run_id
+        print(f"Logging trial {trial.number} data to MLflow run: {target_run_id}")
+
+        # Assign GPU device
+        if gpu_count > 1:
+            gpu_id = trial.number % gpu_count
+            device = torch.device(f"cuda:{gpu_id}")
+            torch.cuda.set_device(device)
+        else:
+            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        return device, mlflow_client, target_run_id
+    
+    def multi_gpu_objective(self, parent_run, trial, epochs, val_interval=5, best_metric="avg_f1", gpu_count=1):
+        """
+        Trains model on multiple GPUs using a parent MLflow run.
+        """
+        self.device, self.client, self.target_run_id = self._setup_parallel_trial_run(trial, parent_run, gpu_count)
+        
+        # Build model + trainer
+        self.my_build_model(trial)
+        self._define_optimizer(trial)
+        model_trainer = trainer.ModelTrainer(
+            self.model, self.device, self.loss_function, 
+            self.metrics_function, self.optimizer
+        )
+        model_trainer.set_parallel_mlflow(self.client, self.target_run_id)
+
+        # Train Model, with error handling
+        score = None
+        run_status = "FAILED"   # default; overwritten on success/prune
+        try:
+            score = self._train_model(
+                trial, model_trainer, epochs, val_interval, 
+                self.sampling['crop_size'], self.sampling['num_samples'], 
+                best_metric)
+            # Save best model and mark run as finished
+            run_status = "FINISHED"
+            self._save_best_model(trial, model_trainer, score)
+            return score
+        except optuna.TrialPruned:
+            run_status = "KILLED"  # communicates early stop 
+            raise
+        except Exception as e:
+            run_status = "FAILED"
+            print(f"[Trial Failed] Unexpected error: {e}")
+            trial.set_user_attr("error", str(e))
+            raise optuna.TrialPruned()
+        finally:
+            self.cleanup(model_trainer)
+            model_trainer.my_log_params({"parent_run_name": parent_run.info.run_name})
+            try:
+                if self.client is not None and self.target_run_id is not None:
+                    self.client.set_terminated(self.target_run_id, status=run_status)
+            except Exception as e:
+                print(f"[Cleanup Failed] Unexpected error: {e}")
+
+    def _save_best_model(self, trial, model_trainer, score):
+        """Saves the best model if it improves upon previous scores."""
+        best_score_so_far = self.get_best_score(trial)
+        if score > best_score_so_far:
+            torch.save(model_trainer.model_weights, f'{self.results_dir}/best_model.pth')
+            io.save_parameters_to_yaml(self.model_builder, model_trainer, self.data_generator, 
+                                    f'{self.results_dir}/model_config.yaml')
+
+    def get_best_score(self, trial):
+        """Retrieve the best score from the trial."""
+        try:
+            return trial.study.best_value
+        except ValueError:
+            return -float('inf')
+
+    def cleanup(self, model_trainer):
+        """Handles cleanup of resources."""
+
+        # Log training parameters
+        params = {
+            'model': self.model_builder.get_model_parameters(),
+            'optimizer': io.get_optimizer_parameters(model_trainer)
+        }    
+        model_trainer.my_log_params(io.flatten_params(params))
+    
+        # Delete the trainer and optimizer objects
+        del model_trainer, self.optimizer
+
+        # If the model object holds GPU memory, delete it explicitly and set it to None
+        if hasattr(self, "model"):
+            del self.model
+            self.model = None
+
+        # Optional: If your model_builder or other objects hold GPU references, delete them too
+        if hasattr(self, "model_builder"):
+            del self.model_builder
+            self.model_builder = None
+
+        # Clear the CUDA cache and force garbage collection
+        torch.cuda.empty_cache()
+        gc.collect()
+

octopi/pytorch/model_search_submitter.py

@@ -0,0 +1,291 @@
+from octopi.datasets import generators, multi_config_generator
+from octopi.utils import config, parsers
+from octopi.pytorch import hyper_search
+import torch, mlflow, optuna
+from typing import List
+import pandas as pd
+
+class ModelSearchSubmit:
+    def __init__(
+        self,
+        copick_config: str,
+        target_name: str,
+        target_user_id: str,
+        target_session_id: str,
+        tomo_algorithm: str,
+        voxel_size: float,
+        model_type: str,
+        best_metric: str = 'avg_f1',
+        num_epochs: int = 1000,
+        num_trials: int = 100,
+        data_split: str = 0.8,
+        random_seed: int = 42,
+        val_interval: int = 10,
+        tomo_batch_size: int = 15,
+        trainRunIDs: List[str] = None,
+        validateRunIDs: List[str] = None,
+        mlflow_experiment_name: str = 'explore',
+    ):
+        """
+        Initialize the ModelSearch class for architecture search with Optuna.
+
+        Args:
+            copick_config (str or dict): Path to the CoPick configuration file or a dictionary for multi-config training.
+            target_name (str): Name of the target for segmentation.
+            target_user_id (str): Optional user ID for tracking.
+            target_session_id (str): Optional session ID for tracking.
+            tomo_algorithm (str): Tomogram algorithm to use.
+            voxel_size (float): Voxel size for tomograms.
+            Nclass (int): Number of prediction classes.
+            model_type (str): Type of model to use.
+            mlflow_experiment_name (str): MLflow experiment name.
+            random_seed (int): Seed for reproducibility.
+            num_epochs (int): Number of epochs per trial.
+            num_trials (int): Number of trials for hyperparameter optimization.
+            tomo_batch_size (int): Batch size for tomogram loading.
+            best_metric (str): Metric to optimize.
+            val_interval (int): Validation interval.
+            trainRunIDs (List[str]): List of training run IDs.
+            validateRunIDs (List[str]): List of validation run IDs.
+            data_split (str): Data split ratios.
+        """
+
+        # Input parameters 
+        self.copick_config = copick_config
+        self.target_name = target_name
+        self.target_user_id = target_user_id
+        self.target_session_id = target_session_id
+        self.tomo_algorithm = tomo_algorithm
+        self.voxel_size = voxel_size
+        self.model_type = model_type
+        self.mlflow_experiment_name = mlflow_experiment_name
+        self.random_seed = random_seed
+        self.num_epochs = num_epochs
+        self.num_trials = num_trials
+        self.tomo_batch_size = tomo_batch_size
+        self.best_metric = best_metric
+        self.val_interval = val_interval
+        self.trainRunIDs = trainRunIDs
+        self.validateRunIDs = validateRunIDs
+        self.data_split = data_split
+        
+        # Data generator - will be initialized in _initialize_data_generator()
+        self.data_generator = None
+
+        # Set random seed for reproducibility
+        config.set_seed(self.random_seed)
+
+        # Initialize dataset generator
+        self._initialize_data_generator()
+
+    def _initialize_data_generator(self):
+        """Initializes the data generator for training and validation datasets."""
+        self._print_input_configs()
+
+        if isinstance(self.copick_config, dict):
+            self.data_generator = multi_config_generator.MultiConfigTrainLoaderManager(
+                self.copick_config,
+                self.target_name,
+                target_session_id=self.target_session_id,
+                target_user_id=self.target_user_id,
+                tomo_algorithm=self.tomo_algorithm,
+                voxel_size=self.voxel_size,
+                tomo_batch_size=self.tomo_batch_size
+            )
+        else:
+            self.data_generator = generators.TrainLoaderManager(
+                self.copick_config,
+                self.target_name,
+                target_session_id=self.target_session_id,
+                target_user_id=self.target_user_id,
+                tomo_algorithm=self.tomo_algorithm,
+                voxel_size=self.voxel_size,
+                tomo_batch_size=self.tomo_batch_size
+            )
+
+        # Split datasets into training and validation
+        ratios = parsers.parse_data_split(self.data_split)
+        self.data_generator.get_data_splits(
+            trainRunIDs=self.trainRunIDs,
+            validateRunIDs=self.validateRunIDs,
+            train_ratio = ratios[0], val_ratio = ratios[1], test_ratio = ratios[2],
+            create_test_dataset = False
+        )
+        
+        # Get the reload frequency
+        self.data_generator.get_reload_frequency(self.num_epochs)
+        self.Nclass = self.data_generator.Nclasses
+        
+    def _print_input_configs(self):
+        """Prints training configuration for debugging purposes."""
+        print(f'\nTraining with:')
+        if isinstance(self.copick_config, dict):
+            for session, config in self.copick_config.items():
+                print(f'  {session}: {config}')
+        else:
+            print(f'  {self.copick_config}')
+        print()
+
+    def run_model_search(self):
+        """Performs model architecture search using Optuna and MLflow."""
+
+        # Set up MLflow tracking
+        try:
+            tracking_uri = config.mlflow_setup()
+            mlflow.set_tracking_uri(tracking_uri)
+        except Exception as e:
+            print(f'Failed to set up MLflow tracking: {e}')
+            pass
+
+        mlflow.set_experiment(self.mlflow_experiment_name)
+    
+        # Create a storage object with heartbeat configuration
+        storage_url = f"sqlite:///explore_results_{self.model_type}/trials.db"
+        self.storage = optuna.storages.RDBStorage(
+            url=storage_url,
+            heartbeat_interval=60,  # Record heartbeat every minute
+            grace_period=600,       # 10 minutes grace period
+            failed_trial_callback=optuna.storages.RetryFailedTrialCallback(max_retry=1)
+        )
+
+        # Detect GPU availability
+        gpu_count = torch.cuda.device_count()
+        print(f'Running Architecture Search Over {gpu_count} GPUs\n')
+
+        # Initialize model search object
+        if gpu_count > 1:
+            self._multi_gpu_optuna(gpu_count)
+        else:
+            model_search = hyper_search.BayesianModelSearch(self.data_generator, self.model_type)
+            self._single_gpu_optuna(model_search)
+
+    def _single_gpu_optuna(self, model_search):
+        """Runs Optuna optimization on a single GPU."""
+        
+        with mlflow.start_run(nested=False) as parent_run:
+
+            model_search.parent_run_id = parent_run.info.run_id
+            model_search.parent_run_name = parent_run.info.run_name
+
+            # Log the experiment parameters
+            mlflow.log_params({"random_seed": self.random_seed})
+            mlflow.log_params(self.data_generator.get_dataloader_parameters())
+            mlflow.log_params({"parent_run_name": parent_run.info.run_name})
+
+            # Determine device
+            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+            # Create the study and run the optimization
+            study = self.get_optuna_study()
+            study.optimize(
+                lambda trial: model_search.objective(
+                    trial, self.num_epochs, device,
+                    val_interval=self.val_interval,
+                    best_metric=self.best_metric,
+                ),
+                n_trials=self.num_trials
+            )
+
+            # Save contour plot
+            self.save_contour_plot_as_png(study)
+
+            print(f"Best trial: {study.best_trial.value}")
+            print(f"Best params: {study.best_params}")
+
+    def _multi_gpu_optuna(self, gpu_count):
+        """Runs Optuna optimization on multiple GPUs."""
+        with mlflow.start_run() as parent_run:
+
+            # Log the experiment parameters
+            mlflow.log_params({"random_seed": self.random_seed})
+            mlflow.log_params(self.data_generator.get_dataloader_parameters())
+            mlflow.log_params({"parent_run_name": parent_run.info.run_name})
+
+            # Run multi-GPU optimization
+            study = self.get_optuna_study()
+            study.optimize(
+                lambda trial: hyper_search.BayesianModelSearch(self.data_generator, self.model_type).multi_gpu_objective(
+                    parent_run, trial,
+                    self.num_epochs,
+                    best_metric=self.best_metric,
+                    val_interval=self.val_interval,
+                    gpu_count=gpu_count
+                ),
+                n_trials=self.num_trials,
+                n_jobs=gpu_count
+            )
+
+            # Save contour Plot
+            self.save_contour_plot_as_png(study)
+
+        print(f"Best trial: {study.best_trial.value}")
+        print(f"Best params: {study.best_params}")
+
+    def _get_optuna_sampler(self):
+        """Returns Optuna's TPE sampler with default settings."""
+        return optuna.samplers.TPESampler(
+            n_startup_trials=10,
+            n_ei_candidates=24,
+            multivariate=True
+        )
+        # return optuna.samplers.BoTorchSampler(
+        #     n_startup_trials=10,
+        #     multivariate=True
+        # )
+
+    def get_optuna_study(self):
+        """Returns the Optuna study object."""
+        return optuna.create_study(
+                storage=self.storage,
+                direction="maximize",
+                sampler=self._get_optuna_sampler(),
+                load_if_exists=True,
+                pruner=self._get_optuna_pruner()
+            )
+
+    def _get_optuna_pruner(self):
+        """Returns Optuna's pruning strategy."""
+        return optuna.pruners.MedianPruner(
+            n_startup_trials=10,    # let at least 10 full trials run before pruning
+            n_warmup_steps=300,     # dont prune before 300 epochs/steps
+            interval_steps=self.val_interval       # check each interval
+        )
+
+    def save_contour_plot_as_png(self, study):
+        """
+        Save the contour plot of hyperparameter interactions as a PNG, 
+        automatically extracting parameter names from the study object.
+
+        Args:
+            study: The Optuna study object.
+            output_path: Path to save the PNG file.
+        """
+        # Extract all parameter names from the study trials
+        all_params = set()
+        for trial in study.trials:
+            all_params.update(trial.params.keys())
+        all_params = list(all_params)  # Convert to a sorted list for consistency
+
+        # Generate the contour plot
+        fig = optuna.visualization.plot_contour(study, params=all_params) 
+
+        # Adjust figure size and font size
+        fig.update_layout(
+            width=6000, height=6000,  # Large figure size
+            font=dict(size=40)  # Increase font size for better readability
+        )
+
+        # Save the plot as a PNG file
+        fig.write_image(f'explore_results_{self.model_type}/contour_plot.png', scale=1)  
+
+        # Extract trial data
+        trials = [
+            {**trial.params, 'objective_value': trial.value}
+            for trial in study.trials if trial.state == optuna.trial.TrialState.COMPLETE
+        ]
+
+        # Convert to DataFrame
+        df = pd.DataFrame(trials)
+
+        # Save to CSV
+        df.to_csv(f"explore_results_{self.model_type}/optuna_results.csv", index=False)