optuna added as base

configs/experiment/catdog_experiment.yaml

@@ -39,7 +39,7 @@ model:
 
 trainer:
   min_epochs: 1
-  max_epochs: 6
+  max_epochs: 5
 
 callbacks:
   model_checkpoint:

configs/trainer/default.yaml

@@ -1,4 +1,3 @@
-_target_: lightning.Trainer
 
 default_root_dir: ${paths.output_dir}
 min_epochs: 1
@@ -10,8 +9,7 @@ devices: auto
 # mixed precision for extra speed-up
 # precision: 16
 
-# set True to to ensure deterministic results
-# makes training slower but gives more reproducibility than just setting seeds
+# set True to to ensure deterministic results makes training slower but gives more reproducibility than just setting seeds
 deterministic: True
 
 # Log every N steps in training and validation
@@ -19,3 +17,4 @@ log_every_n_steps: 10
 fast_dev_run: False
 
 gradient_clip_val: 1.0
+gradient_clip_algorithm: 'norm'

src/train_new.py

@@ -1,7 +1,5 @@
 """
-Train and evaluate a model using PyTorch Lightning.
-Initializes the DataModule, Model, Trainer, and runs training and testing.
-Initializes loggers and callbacks from the configuration using Hydra and target paths from the configuration.
+Train and evaluate a model using PyTorch Lightning with Optuna for hyperparameter optimization.
 """
 
 import os
@@ -17,51 +15,34 @@ from src.utils.logging_utils import setup_logger, task_wrapper
 from loguru import logger
 import rootutils
 from lightning.pytorch.loggers import Logger
-from lightning.pytorch.callbacks import Callback
+import optuna
+from lightning.pytorch import Trainer
 
 # Load environment variables
 load_dotenv(find_dotenv(".env"))
 
 # Setup root directory
-
 root = rootutils.setup_root(__file__, indicator=".project-root")
 
 
-def instantiate_callbacks(callback_cfg: DictConfig) -> List[Callback]:
-    """Instantiate and return a list of callbacks from the configuration."""
-    callbacks_ls: List[L.Callback] = []
-
-    if not callback_cfg:
-        logger.warning("No callback configs found! Skipping..")
-        return None
-
-    if not isinstance(callback_cfg, DictConfig):
-        raise TypeError("Callbacks config must be a DictConfig!")
-
-    for _, cb_conf in callback_cfg.items():
-        if "_target_" in cb_conf:
-            logger.info(f"Instantiating callback <{cb_conf._target_}>")
-            callbacks_ls.append(hydra.utils.instantiate(cb_conf))
-
-    return callbacks_ls
-
-
 def instantiate_loggers(logger_cfg: DictConfig) -> List[Logger]:
     """Instantiate and return a list of loggers from the configuration."""
     loggers_ls: List[Logger] = []
 
-    if not logger_cfg:
-        logger.warning("No logger configs found! Skipping..")
+    if not logger_cfg or isinstance(logger_cfg, bool):
+        logger.warning("No valid logger configs found! Skipping..")
         return loggers_ls
 
     if not isinstance(logger_cfg, DictConfig):
         raise TypeError("Logger config must be a DictConfig!")
 
     for _, lg_conf in logger_cfg.items():
-        if "_target_" in lg_conf:
+        if isinstance(lg_conf, DictConfig) and "_target_" in lg_conf:
             logger.info(f"Instantiating logger <{lg_conf._target_}>")
-            loggers_ls.append(hydra.utils.instantiate(lg_conf))
-
+            try:
+                loggers_ls.append(hydra.utils.instantiate(lg_conf))
+            except Exception as e:
+                logger.error(f"Failed to instantiate logger {lg_conf}: {e}")
     return loggers_ls
 
 
@@ -93,16 +74,19 @@ def clear_checkpoint_directory(ckpt_dir: str):
 def train_module(
     data_module: L.LightningDataModule, model: L.LightningModule, trainer: L.Trainer
 ):
-    """Train the model and log metrics."""
-    logger.info("Starting training")
+    """Train the model, return validation accuracy for each epoch."""
+    logger.info("Starting training with custom pruning")
+
     trainer.fit(model, data_module)
-    train_metrics = trainer.callback_metrics
-    train_acc = train_metrics.get("train_acc")
-    val_acc = train_metrics.get("val_acc")
-    logger.info(
-        f"Training completed. Metrics - train_acc: {train_acc}, val_acc: {val_acc}"
-    )
-    return train_metrics
+    val_accuracies = []
+
+    for epoch in range(trainer.current_epoch):
+        val_acc = trainer.callback_metrics.get("val_acc")
+        if val_acc is not None:
+            val_accuracies.append(val_acc.item())
+            logger.info(f"Epoch {epoch}: val_acc={val_acc}")
+
+    return val_accuracies
 
 
 @task_wrapper
@@ -122,77 +106,78 @@ def run_test_module(
     return test_metrics[0] if test_metrics else {}
 
 
-@hydra.main(config_path="../configs", config_name="train", version_base="1.3")
-def setup_run_trainer(cfg: DictConfig):
-    """Set up and run the Trainer for training and testing."""
-    # Display configuration
-    logger.info(f"Config:\n{OmegaConf.to_yaml(cfg)}")
+def objective(trial: optuna.trial.Trial, cfg: DictConfig):
+    """Objective function for Optuna hyperparameter tuning."""
 
-    # Initialize logger
-    log_path = Path(cfg.paths.log_dir) / (
-        "train.log" if cfg.task_name == "train" else "eval.log"
-    )
-    setup_logger(log_path)
-
-    # Display key paths
-    for path_name in [
-        "root_dir",
-        "data_dir",
-        "log_dir",
-        "ckpt_dir",
-        "artifact_dir",
-        "output_dir",
-    ]:
-        logger.info(
-            f"{path_name.replace('_', ' ').capitalize()}: {cfg.paths[path_name]}"
-        )
+    # Sample hyperparameters for the model
+    cfg.model.embed_dim = trial.suggest_categorical("embed_dim", [64, 128, 256])
+    cfg.model.depth = trial.suggest_int("depth", 2, 6)
+    cfg.model.lr = trial.suggest_loguniform("lr", 1e-5, 1e-3)
+    cfg.model.mlp_ratio = trial.suggest_float("mlp_ratio", 1.0, 4.0)
 
-    # Initialize DataModule and Model
-    logger.info(f"Instantiating datamodule <{cfg.data._target_}>")
-    datamodule: L.LightningDataModule = hydra.utils.instantiate(cfg.data)
-    logger.info(f"Instantiating model <{cfg.model._target_}>")
+    # Initialize data module and model
+    data_module: L.LightningDataModule = hydra.utils.instantiate(cfg.data)
     model: L.LightningModule = hydra.utils.instantiate(cfg.model)
 
-    # Check GPU availability and set seed for reproducibility
-    logger.info("GPU available" if torch.cuda.is_available() else "No GPU available")
-    L.seed_everything(cfg.seed, workers=True)
-
-    # Set up callbacks, loggers, and Trainer
-    callbacks = instantiate_callbacks(cfg.callbacks)
-    logger.info(f"Callbacks: {callbacks}")
+    # Set up logger
     loggers = instantiate_loggers(cfg.logger)
-    logger.info(f"Loggers: {loggers}")
-    trainer: L.Trainer = hydra.utils.instantiate(
-        cfg.trainer, callbacks=callbacks, logger=loggers
+
+    # Trainer configuration without pruning callback
+    trainer = Trainer(**cfg.trainer, logger=loggers)
+
+    # Clear checkpoint directory
+    clear_checkpoint_directory(cfg.paths.ckpt_dir)
+
+    # Train and get val_acc for each epoch
+    val_accuracies = train_module(data_module, model, trainer)
+
+    # Report validation accuracy and prune if necessary
+    for epoch, val_acc in enumerate(val_accuracies):
+        trial.report(val_acc, step=epoch)
+
+        # Check if the trial should be pruned at this epoch
+        if trial.should_prune():
+            logger.info(f"Pruning trial at epoch {epoch}")
+            raise optuna.TrialPruned()
+
+    # Return the final validation accuracy as the objective metric
+    return val_accuracies[-1] if val_accuracies else 0.0
+
+
+@hydra.main(config_path="../configs", config_name="train", version_base="1.3")
+def setup_trainer(cfg: DictConfig):
+    logger.info(f"Config:\n{OmegaConf.to_yaml(cfg)}")
+
+    setup_logger(
+        Path(cfg.paths.log_dir)
+        / ("train.log" if cfg.task_name == "train" else "eval.log")
     )
 
-    # Training phase
-    train_metrics = {}
-    if cfg.get("train"):
-        clear_checkpoint_directory(cfg.paths.ckpt_dir)
-        train_metrics = train_module(datamodule, model, trainer)
-        (Path(cfg.paths.ckpt_dir) / "train_done.flag").write_text(
-            "Training completed.\n"
+    if cfg.get("train", False):
+        pruner = optuna.pruners.MedianPruner()
+        study = optuna.create_study(
+            direction="maximize", pruner=pruner, study_name="pytorch_lightning_optuna"
         )
-
-    # Testing phase
-    test_metrics = {}
-    if cfg.get("test"):
-        test_metrics = run_test_module(cfg, datamodule, model, trainer)
-
-    # Combine metrics and extract optimization metric
-    all_metrics = {**train_metrics, **test_metrics}
-    optimization_metric = all_metrics.get(cfg.get("optimization_metric"), 0.0)
-    (
-        logger.warning(
-            f"Optimization metric '{cfg.get('optimization_metric')}' not found. Defaulting to 0."
+        study.optimize(
+            lambda trial: objective(trial, cfg), n_trials=3, show_progress_bar=True
         )
-        if optimization_metric == 0.0
-        else logger.info(f"Optimization metric: {optimization_metric}")
-    )
 
-    return optimization_metric
+        # Log best trial results
+        best_trial = study.best_trial
+        logger.info(f"Best trial number: {best_trial.number}")
+        logger.info(f"Best trial value (val_acc): {best_trial.value}")
+        for key, value in best_trial.params.items():
+            logger.info(f"  {key}: {value}")
+
+    if cfg.get("test", False):
+        data_module: L.LightningDataModule = hydra.utils.instantiate(cfg.data)
+        model: L.LightningModule = hydra.utils.instantiate(cfg.model)
+        trainer = Trainer(**cfg.trainer, logger=instantiate_loggers(cfg.logger))
+        test_metrics = run_test_module(cfg, data_module, model, trainer)
+        logger.info(f"Test metrics: {test_metrics}")
+
+    return cfg.model if not cfg.get("test", False) else test_metrics
 
 
 if __name__ == "__main__":
-    setup_run_trainer()
+    setup_trainer()

src/train_old.py

@@ -0,0 +1,260 @@
+import os
+import shutil
+from pathlib import Path
+import torch
+import lightning as L
+from lightning.pytorch.loggers import Logger
+from typing import List
+from src.datamodules.dogbreed_datamodule import main_dataloader
+from src.utils.logging_utils import setup_logger, task_wrapper
+from loguru import logger
+from dotenv import load_dotenv, find_dotenv
+import rootutils
+import hydra
+from omegaconf import DictConfig, OmegaConf
+
+# Load environment variables
+load_dotenv(find_dotenv(".env"))
+
+# Setup root directory
+root = rootutils.setup_root(__file__, indicator=".project-root")
+
+
+def instantiate_callbacks(callback_cfg: DictConfig) -> List[L.Callback]:
+    """Instantiate and return a list of callbacks from the configuration."""
+    callbacks: List[L.Callback] = []
+
+    if not callback_cfg:
+        logger.warning("No callback configs found! Skipping..")
+        return callbacks
+
+    if not isinstance(callback_cfg, DictConfig):
+        raise TypeError("Callbacks config must be a DictConfig!")
+
+    for _, cb_conf in callback_cfg.items():
+        if "_target_" in cb_conf:
+            logger.info(f"Instantiating callback <{cb_conf._target_}>")
+            callbacks.append(hydra.utils.instantiate(cb_conf))
+
+    return callbacks
+
+
+def instantiate_loggers(logger_cfg: DictConfig) -> List[Logger]:
+    """Instantiate and return a list of loggers from the configuration."""
+    loggers_ls: List[Logger] = []
+
+    if not logger_cfg:
+        logger.warning("No logger configs found! Skipping..")
+        return loggers_ls
+
+    if not isinstance(logger_cfg, DictConfig):
+        raise TypeError("Logger config must be a DictConfig!")
+
+    for _, lg_conf in logger_cfg.items():
+        if "_target_" in lg_conf:
+            logger.info(f"Instantiating logger <{lg_conf._target_}>")
+            loggers_ls.append(hydra.utils.instantiate(lg_conf))
+
+    return loggers_ls
+
+
+def load_checkpoint_if_available(ckpt_path: str) -> str:
+    """Check if the specified checkpoint exists and return the valid checkpoint path."""
+    if ckpt_path and Path(ckpt_path).exists():
+        logger.info(f"Checkpoint found: {ckpt_path}")
+        return ckpt_path
+    else:
+        logger.warning(
+            f"No checkpoint found at {ckpt_path}. Using current model weights."
+        )
+        return None
+
+
+def clear_checkpoint_directory(ckpt_dir: str):
+    """Clear all contents of the checkpoint directory without deleting the directory itself."""
+    ckpt_dir_path = Path(ckpt_dir)
+    if ckpt_dir_path.exists() and ckpt_dir_path.is_dir():
+        logger.info(f"Clearing checkpoint directory: {ckpt_dir}")
+        # Iterate over all files and directories in the checkpoint directory and remove them
+        for item in ckpt_dir_path.iterdir():
+            try:
+                if item.is_file() or item.is_symlink():
+                    item.unlink()  # Remove file or symlink
+                elif item.is_dir():
+                    shutil.rmtree(item)  # Remove directory
+            except Exception as e:
+                logger.error(f"Failed to delete {item}: {e}")
+        logger.info(f"Checkpoint directory cleared: {ckpt_dir}")
+    else:
+        logger.info(
+            f"Checkpoint directory does not exist. Creating directory: {ckpt_dir}"
+        )
+        os.makedirs(ckpt_dir_path, exist_ok=True)
+
+
+@task_wrapper
+def train_module(
+    cfg: DictConfig,
+    data_module: L.LightningDataModule,
+    model: L.LightningModule,
+    trainer: L.Trainer,
+):
+    """Train the model using the provided Trainer and DataModule."""
+    logger.info("Training the model")
+    trainer.fit(model, data_module)
+    train_metrics = trainer.callback_metrics
+    try:
+        logger.info(
+            f"Training completed with the following metrics- train_acc: {train_metrics['train_acc'].item()} and val_acc: {train_metrics['val_acc'].item()}"
+        )
+    except KeyError:
+        logger.info(f"Training completed with the following metrics:{train_metrics}")
+
+    return train_metrics
+
+
+@task_wrapper
+def run_test_module(
+    cfg: DictConfig,
+    datamodule: L.LightningDataModule,
+    model: L.LightningModule,
+    trainer: L.Trainer,
+):
+    """Test the model using the best checkpoint or the current model weights."""
+    logger.info("Testing the model")
+    datamodule.setup(stage="test")
+
+    ckpt_path = load_checkpoint_if_available(cfg.ckpt_path)
+
+    # If no checkpoint is available, Lightning will use current model weights
+    test_metrics = trainer.test(model, datamodule, ckpt_path=ckpt_path)
+    logger.info(f"Test metrics:\n{test_metrics}")
+
+    return test_metrics[0] if test_metrics else {}
+
+
+@hydra.main(config_path="../configs", config_name="train", version_base="1.1")
+def setup_run_trainer(cfg: DictConfig):
+    """Set up and run the Trainer for training and testing the model."""
+    # show me the entire config
+    logger.info(f"Config:\n{OmegaConf.to_yaml(cfg)}")
+    # Initialize logger
+    if cfg.task_name == "train":
+        log_path = Path(cfg.paths.log_dir) / "train.log"
+    else:
+        log_path = Path(cfg.paths.log_dir) / "eval.log"
+    setup_logger(log_path)
+
+    # the path to the checkpoint directory
+    root_dir = cfg.paths.root_dir
+    logger.info(f"Root directory: {root_dir}")
+
+    logger.info(f"Current working directory: {os.listdir(root_dir)}")
+
+    ckpt_dir = cfg.paths.ckpt_dir
+    logger.info(f"Checkpoint directory: {ckpt_dir}")
+
+    # the path to the data directory
+    data_dir = cfg.paths.data_dir
+    logger.info(f"Data directory: {data_dir}")
+
+    # the path to the log directory
+    log_dir = cfg.paths.log_dir
+    logger.info(f"Log directory: {log_dir}")
+
+    # the path to the artifact directory
+    artifact_dir = cfg.paths.artifact_dir
+    logger.info(f"Artifact directory: {artifact_dir}")
+
+    # output directory
+    output_dir = cfg.paths.output_dir
+    logger.info(f"Output directory: {output_dir}")
+
+    # name of the experiment
+    experiment_name = cfg.name
+    logger.info(f"Experiment name: {experiment_name}")
+
+    # Initialize DataModule
+    if experiment_name == "dogbreed_experiment":
+        logger.info("Setting up the DataModule")
+        dataset_df, datamodule = main_dataloader(cfg)
+        labels = dataset_df.label.nunique()
+        logger.info(f"Number of classes: {labels}")
+
+        os.makedirs(cfg.paths.artifact_dir, exist_ok=True)
+        dataset_df.to_csv(
+            Path(cfg.paths.artifact_dir) / "dogbreed_dataset.csv", index=False
+        )
+    elif (
+        experiment_name == "catdog_experiment"
+        or experiment_name == "catdog_experiment_convnext"
+    ):
+        # Initialize DataModule
+        logger.info(f"Instantiating datamodule <{cfg.data._target_}>")
+        datamodule: L.LightningDataModule = hydra.utils.instantiate(cfg.data)
+
+    # Check for GPU availability
+    logger.info("GPU available" if torch.cuda.is_available() else "No GPU available")
+
+    # Set seed for reproducibility
+    L.seed_everything(cfg.seed, workers=True)
+
+    # Initialize model
+    logger.info(f"Instantiating model <{cfg.model._target_}>")
+    model: L.LightningModule = hydra.utils.instantiate(cfg.model)
+
+    logger.info(f"Model summary:\n{model}")
+
+    # Set up callbacks and loggers
+    logger.info("Setting up callbacks and loggers")
+    callbacks: List[L.Callback] = instantiate_callbacks(cfg.get("callbacks"))
+    logger.info(f"Callbacks: {callbacks}")
+    loggers: List[Logger] = instantiate_loggers(cfg.get("logger"))
+    logger.info(f"Loggers: {loggers}")
+
+    # Initialize Trainer
+    logger.info(f"Instantiating trainer <{cfg.trainer._target_}>")
+    trainer: L.Trainer = hydra.utils.instantiate(
+        cfg.trainer, callbacks=callbacks, logger=loggers
+    )
+
+    # Train and test the model based on config settings
+    train_metrics = {}
+    if cfg.get("train"):
+        # clear the checkpoint directory
+        clear_checkpoint_directory(cfg.paths.ckpt_dir)
+
+        logger.info("Training the model")
+        train_metrics = train_module(cfg, datamodule, model, trainer)
+
+        # Write training done flag using Hydra paths config
+        done_flag_path = Path(cfg.paths.ckpt_dir) / "train_done.flag"
+        with done_flag_path.open("w") as f:
+            f.write("Training completed.\n")
+        logger.info(f"Training completion flag written to: {done_flag_path}")
+
+        logger.info(
+            f"Training completed. Checkpoint directory: {os.listdir(cfg.paths.ckpt_dir)}"
+        )
+
+    test_metrics = {}
+    if cfg.get("test"):
+        logger.info(f"Checkpoint directory: {os.listdir(cfg.paths.ckpt_dir)}")
+        test_metrics = run_test_module(cfg, datamodule, model, trainer)
+
+    # Combine metrics
+    all_metrics = {**train_metrics, **test_metrics}
+
+    # Extract and return the optimization metric
+    optimization_metric = all_metrics.get(cfg.get("optimization_metric"))
+    if optimization_metric is None:
+        logger.warning(
+            f"Optimization metric '{cfg.get('optimization_metric')}' not found in metrics. Returning 0."
+        )
+        return 0.0
+
+    return optimization_metric
+
+
+if __name__ == "__main__":
+    setup_run_trainer()

src/{train.py → train_optuna_callbacks.py}

@@ -1,7 +1,5 @@
 """
-Train and evaluate a model using PyTorch Lightning.
-Initializes the DataModule, Model, Trainer, and runs training and testing.
-Initializes loggers and callbacks from the configuration using Hydra configuration but with a more modular approach without direct instantiation.
+Train and evaluate a model using PyTorch Lightning with Optuna for hyperparameter optimization.
 """
 
 import os
@@ -10,47 +8,61 @@ from pathlib import Path
 from typing import List
 import torch
 import lightning as L
-from lightning.pytorch.loggers import Logger, TensorBoardLogger, CSVLogger
-from lightning.pytorch.callbacks import (
-    ModelCheckpoint,
-    EarlyStopping,
-    RichModelSummary,
-    RichProgressBar,
-)
 from dotenv import load_dotenv, find_dotenv
 import hydra
 from omegaconf import DictConfig, OmegaConf
-from src.datamodules.catdog_datamodule import CatDogImageDataModule
 from src.utils.logging_utils import setup_logger, task_wrapper
 from loguru import logger
 import rootutils
+from lightning.pytorch.loggers import Logger
+import optuna
+from lightning.pytorch import Trainer
 
 # Load environment variables
 load_dotenv(find_dotenv(".env"))
 
 # Setup root directory
-
 root = rootutils.setup_root(__file__, indicator=".project-root")
 
 
-def initialize_callbacks(cfg: DictConfig) -> List[L.Callback]:
-    """Initialize callbacks based on configuration."""
-    callback_classes = {
-        "model_checkpoint": ModelCheckpoint,
-        "early_stopping": EarlyStopping,
-        "rich_model_summary": RichModelSummary,
-        "rich_progress_bar": RichProgressBar,
-    }
-    return [callback_classes[name](**params) for name, params in cfg.callbacks.items()]
+def instantiate_callbacks(callback_cfg: DictConfig) -> List[L.Callback]:
+    """Instantiate and return a list of callbacks from the configuration."""
+    callbacks: List[L.Callback] = []
+
+    if not callback_cfg:
+        logger.warning("No callback configs found! Skipping..")
+        return callbacks
+
+    if not isinstance(callback_cfg, DictConfig):
+        raise TypeError("Callbacks config must be a DictConfig!")
+
+    for _, cb_conf in callback_cfg.items():
+        if "_target_" in cb_conf:
+            logger.info(f"Instantiating callback <{cb_conf._target_}>")
+            callbacks.append(hydra.utils.instantiate(cb_conf))
+
+    return callbacks
 
 
-def initialize_loggers(cfg: DictConfig) -> List[Logger]:
-    """Initialize loggers based on configuration."""
-    logger_classes = {
-        "tensorboard": TensorBoardLogger,
-        "csv": CSVLogger,
-    }
-    return [logger_classes[name](**params) for name, params in cfg.logger.items()]
+def instantiate_loggers(logger_cfg: DictConfig) -> List[Logger]:
+    """Instantiate and return a list of loggers from the configuration."""
+    loggers_ls: List[Logger] = []
+
+    if not logger_cfg or isinstance(logger_cfg, bool):
+        logger.warning("No valid logger configs found! Skipping..")
+        return loggers_ls
+
+    if not isinstance(logger_cfg, DictConfig):
+        raise TypeError("Logger config must be a DictConfig!")
+
+    for _, lg_conf in logger_cfg.items():
+        if isinstance(lg_conf, DictConfig) and "_target_" in lg_conf:
+            logger.info(f"Instantiating logger <{lg_conf._target_}>")
+            try:
+                loggers_ls.append(hydra.utils.instantiate(lg_conf))
+            except Exception as e:
+                logger.error(f"Failed to instantiate logger {lg_conf}: {e}")
+    return loggers_ls
 
 
 def load_checkpoint_if_available(ckpt_path: str) -> str:
@@ -81,16 +93,19 @@ def clear_checkpoint_directory(ckpt_dir: str):
 def train_module(
     data_module: L.LightningDataModule, model: L.LightningModule, trainer: L.Trainer
 ):
-    """Train the model and log metrics."""
-    logger.info("Starting training")
+    """Train the model, return validation accuracy for each epoch."""
+    logger.info("Starting training with custom pruning")
+
     trainer.fit(model, data_module)
-    train_metrics = trainer.callback_metrics
-    train_acc = train_metrics.get("train_acc")
-    val_acc = train_metrics.get("val_acc")
-    logger.info(
-        f"Training completed. Metrics - train_acc: {train_acc}, val_acc: {val_acc}"
-    )
-    return train_metrics
+    val_accuracies = []
+
+    for epoch in range(trainer.current_epoch):
+        val_acc = trainer.callback_metrics.get("val_acc")
+        if val_acc is not None:
+            val_accuracies.append(val_acc.item())
+            logger.info(f"Epoch {epoch}: val_acc={val_acc}")
+
+    return val_accuracies
 
 
 @task_wrapper
@@ -110,77 +125,84 @@ def run_test_module(
     return test_metrics[0] if test_metrics else {}
 
 
-@hydra.main(config_path="../configs", config_name="train", version_base="1.1")
-def setup_run_trainer(cfg: DictConfig):
-    """Set up and run the Trainer for training and testing."""
-    # Display configuration
-    logger.info(f"Config:\n{OmegaConf.to_yaml(cfg)}")
+def objective(trial: optuna.trial.Trial, cfg: DictConfig, callbacks: List[L.Callback]):
+    """Objective function for Optuna hyperparameter tuning."""
 
-    # Initialize logger
-    log_path = Path(cfg.paths.log_dir) / (
-        "train.log" if cfg.task_name == "train" else "eval.log"
-    )
-    setup_logger(log_path)
-
-    # Display key paths
-    for path_name in [
-        "root_dir",
-        "data_dir",
-        "log_dir",
-        "ckpt_dir",
-        "artifact_dir",
-        "output_dir",
-    ]:
-        logger.info(
-            f"{path_name.replace('_', ' ').capitalize()}: {cfg.paths[path_name]}"
-        )
+    # Sample hyperparameters for the model
+    cfg.model.embed_dim = trial.suggest_categorical("embed_dim", [64, 128, 256])
+    cfg.model.depth = trial.suggest_int("depth", 2, 6)
+    cfg.model.lr = trial.suggest_loguniform("lr", 1e-5, 1e-3)
+    cfg.model.mlp_ratio = trial.suggest_float("mlp_ratio", 1.0, 4.0)
 
-    # Initialize DataModule and Model
-    logger.info(f"Instantiating datamodule <{cfg.data._target_}>")
-    datamodule: L.LightningDataModule = hydra.utils.instantiate(cfg.data)
-    logger.info(f"Instantiating model <{cfg.model._target_}>")
+    # Initialize data module and model
+    data_module: L.LightningDataModule = hydra.utils.instantiate(cfg.data)
     model: L.LightningModule = hydra.utils.instantiate(cfg.model)
 
-    # Check GPU availability and set seed for reproducibility
-    logger.info("GPU available" if torch.cuda.is_available() else "No GPU available")
-    L.seed_everything(cfg.seed, workers=True)
+    # Set up logger
+    loggers = instantiate_loggers(cfg.logger)
 
-    # Set up callbacks, loggers, and Trainer
-    callbacks = initialize_callbacks(cfg)
-    logger.info(f"Callbacks: {callbacks}")
-    loggers = initialize_loggers(cfg)
-    logger.info(f"Loggers: {loggers}")
-    trainer: L.Trainer = hydra.utils.instantiate(
-        cfg.trainer, callbacks=callbacks, logger=loggers
+    # Trainer configuration with passed callbacks
+    trainer = Trainer(**cfg.trainer, logger=loggers, callbacks=callbacks)
+
+    # Clear checkpoint directory
+    clear_checkpoint_directory(cfg.paths.ckpt_dir)
+
+    # Train and get val_acc for each epoch
+    val_accuracies = train_module(data_module, model, trainer)
+
+    # Report validation accuracy and prune if necessary
+    for epoch, val_acc in enumerate(val_accuracies):
+        trial.report(val_acc, step=epoch)
+
+        # Check if the trial should be pruned at this epoch
+        if trial.should_prune():
+            logger.info(f"Pruning trial at epoch {epoch}")
+            raise optuna.TrialPruned()
+
+    # Return the final validation accuracy as the objective metric
+    return val_accuracies[-1] if val_accuracies else 0.0
+
+
+@hydra.main(config_path="../configs", config_name="train", version_base="1.3")
+def setup_trainer(cfg: DictConfig):
+    logger.info(f"Config:\n{OmegaConf.to_yaml(cfg)}")
+
+    setup_logger(
+        Path(cfg.paths.log_dir)
+        / ("train.log" if cfg.task_name == "train" else "eval.log")
     )
 
-    # Training phase
-    train_metrics = {}
-    if cfg.get("train"):
-        clear_checkpoint_directory(cfg.paths.ckpt_dir)
-        train_metrics = train_module(datamodule, model, trainer)
-        (Path(cfg.paths.ckpt_dir) / "train_done.flag").write_text(
-            "Training completed.\n"
-        )
+    # Instantiate callbacks
+    callbacks = instantiate_callbacks(cfg.callbacks)
+    logger.info(f"Callbacks: {callbacks}")
 
-    # Testing phase
-    test_metrics = {}
-    if cfg.get("test"):
-        test_metrics = run_test_module(cfg, datamodule, model, trainer)
-
-    # Combine metrics and extract optimization metric
-    all_metrics = {**train_metrics, **test_metrics}
-    optimization_metric = all_metrics.get(cfg.get("optimization_metric"), 0.0)
-    (
-        logger.warning(
-            f"Optimization metric '{cfg.get('optimization_metric')}' not found. Defaulting to 0."
+    if cfg.get("train", False):
+        pruner = optuna.pruners.MedianPruner()
+        study = optuna.create_study(
+            direction="maximize", pruner=pruner, study_name="pytorch_lightning_optuna"
         )
-        if optimization_metric == 0.0
-        else logger.info(f"Optimization metric: {optimization_metric}")
-    )
+        study.optimize(
+            lambda trial: objective(trial, cfg, callbacks),
+            n_trials=5,
+            show_progress_bar=True,
+        )
+
+        # Log best trial results
+        best_trial = study.best_trial
+        logger.info(f"Best trial number: {best_trial.number}")
+        logger.info(f"Best trial value (val_acc): {best_trial.value}")
+        for key, value in best_trial.params.items():
+            logger.info(f"  {key}: {value}")
+
+    if cfg.get("test", False):
+        data_module: L.LightningDataModule = hydra.utils.instantiate(cfg.data)
+        model: L.LightningModule = hydra.utils.instantiate(cfg.model)
+        trainer = Trainer(**cfg.trainer, logger=instantiate_loggers(cfg.logger))
+        test_metrics = run_test_module(cfg, data_module, model, trainer)
+        logger.info(f"Test metrics: {test_metrics}")
 
-    return optimization_metric
+    return cfg.model if not cfg.get("test", False) else test_metrics
 
 
 if __name__ == "__main__":
-    setup_run_trainer()
+    setup_trainer()