qadence 1.8.0__py3-none-any.whl → 1.9.1__py3-none-any.whl

qadence/ml_tools/{saveload.py → callbacks/saveload.py}

@@ -11,7 +11,7 @@ from nevergrad.optimization.base import Optimizer as NGOptimizer
 from torch.nn import Module
 from torch.optim import Optimizer
 
-logger = getLogger(__name__)
+logger = getLogger("ml_tools")
 
 
 def get_latest_checkpoint_name(folder: Path, type: str, device: str | torch.device = "cpu") -> Path:
@@ -19,6 +19,7 @@ def get_latest_checkpoint_name(folder: Path, type: str, device: str | torch.devi
     files = [f for f in os.listdir(folder) if f.endswith(".pt") and type in f]
     if len(files) == 0:
         logger.error(f"Directory {folder} does not contain any {type} checkpoints.")
+        pass
     if len(files) == 1:
         file = Path(files[0])
     else:
@@ -66,8 +67,7 @@ def write_checkpoint(
     iteration: int | str,
 ) -> None:
     from qadence import QuantumModel
-
-    from .models import QNN
+    from qadence.ml_tools.models import QNN
 
     device = None
     try:
@@ -79,10 +79,8 @@ def write_checkpoint(
         )
         device = str(device).split(":")[0]  # in case of using several CUDA devices
     except Exception as e:
-        msg = (
-            f"Unable to identify in which device the QuantumModel is stored due to {e}."
-            "Setting device to None"
-        )
+        msg = f"""Unable to identify in which device the QuantumModel is stored due to {e}.
+                            Setting device to None"""
         logger.warning(msg)
 
     iteration_substring = f"{iteration:03n}" if isinstance(iteration, int) else iteration
@@ -135,7 +133,9 @@ def load_model(
         model_ckpt_name = get_latest_checkpoint_name(folder, "model", device)
 
     try:
-        iteration, model_dict = torch.load(folder / model_ckpt_name, *args, **kwargs)
+        iteration, model_dict = torch.load(
+            folder / model_ckpt_name, weights_only=False, *args, **kwargs
+        )
         if isinstance(model, (QuantumModel, QNN)):
             model.load_params_from_dict(model_dict)
         elif isinstance(model, Module):
@@ -146,8 +146,8 @@ def load_model(
             model.to(device)
 
     except Exception as e:
-        msg = f"Unable to load state dict due to {e}.\
-               No corresponding pre-trained model found. Returning the un-trained model."
+        msg = f"""Unable to load state dict due to {e}.
+                            No corresponding pre-trained model found."""
         logger.warning(msg)
     return model, iteration
 
@@ -162,7 +162,7 @@ def load_optimizer(
         opt_ckpt_name = get_latest_checkpoint_name(folder, "opt", device)
     if os.path.isfile(folder / opt_ckpt_name):
         if isinstance(optimizer, Optimizer):
-            (_, OptType, optimizer_state) = torch.load(folder / opt_ckpt_name)
+            (_, OptType, optimizer_state) = torch.load(folder / opt_ckpt_name, weights_only=False)
             if isinstance(optimizer, OptType):
                 optimizer.load_state_dict(optimizer_state)
 

qadence/ml_tools/callbacks/writer_registry.py

@@ -0,0 +1,441 @@
+from __future__ import annotations
+
+import os
+from abc import ABC, abstractmethod
+from logging import getLogger
+from types import ModuleType
+from typing import Any, Callable, Union
+from uuid import uuid4
+
+from matplotlib.figure import Figure
+from torch import Tensor
+from torch.nn import Module
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+
+from qadence.ml_tools.config import TrainConfig
+from qadence.ml_tools.data import DictDataLoader, OptimizeResult
+from qadence.types import ExperimentTrackingTool
+
+logger = getLogger("ml_tools")
+
+# Type aliases
+PlottingFunction = Callable[[Module, int], tuple[str, Figure]]
+InputData = Union[Tensor, dict[str, Tensor]]
+
+
+class BaseWriter(ABC):
+    """
+    Abstract base class for experiment tracking writers.
+
+    Methods:
+        open(config, iteration=None): Opens the writer and sets up the logging
+            environment.
+        close(): Closes the writer and finalizes any ongoing logging processes.
+        print_metrics(result): Prints metrics and loss in a formatted manner.
+        write(result): Writes the optimization results to the tracking tool.
+        log_hyperparams(hyperparams): Logs the hyperparameters to the tracking tool.
+        plot(model, iteration, plotting_functions): Logs model plots using provided
+            plotting functions.
+        log_model(model, dataloader): Logs the model and any relevant information.
+    """
+
+    run: Any  # [attr-defined]
+
+    @abstractmethod
+    def open(self, config: TrainConfig, iteration: int | None = None) -> Any:
+        """
+        Opens the writer and prepares it for logging.
+
+        Args:
+            config: Configuration object containing settings for logging.
+            iteration (int, optional): The iteration step to start logging from.
+                Defaults to None.
+        """
+        raise NotImplementedError("Writers must implement an open method.")
+
+    @abstractmethod
+    def close(self) -> None:
+        """Closes the writer and finalizes logging."""
+        raise NotImplementedError("Writers must implement a close method.")
+
+    @abstractmethod
+    def write(self, result: OptimizeResult) -> None:
+        """
+        Logs the results of the current iteration.
+
+        Args:
+            result (OptimizeResult): The optimization results to log.
+        """
+        raise NotImplementedError("Writers must implement a write method.")
+
+    @abstractmethod
+    def log_hyperparams(self, hyperparams: dict) -> None:
+        """
+        Logs hyperparameters.
+
+        Args:
+            hyperparams (dict): A dictionary of hyperparameters to log.
+        """
+        raise NotImplementedError("Writers must implement a log_hyperparams method.")
+
+    @abstractmethod
+    def plot(
+        self,
+        model: Module,
+        iteration: int,
+        plotting_functions: tuple[PlottingFunction, ...],
+    ) -> None:
+        """
+        Logs plots of the model using provided plotting functions.
+
+        Args:
+            model (Module): The model to plot.
+            iteration (int): The current iteration number.
+            plotting_functions (tuple[PlottingFunction, ...]): Functions used to
+                generate plots.
+        """
+        raise NotImplementedError("Writers must implement a plot method.")
+
+    @abstractmethod
+    def log_model(
+        self,
+        model: Module,
+        train_dataloader: DataLoader | DictDataLoader | None = None,
+        val_dataloader: DataLoader | DictDataLoader | None = None,
+        test_dataloader: DataLoader | DictDataLoader | None = None,
+    ) -> None:
+        """
+        Logs the model and associated data.
+
+        Args:
+            model (Module): The model to log.
+            train_dataloader (DataLoader | DictDataLoader |  None): DataLoader for training data.
+            val_dataloader (DataLoader | DictDataLoader |  None): DataLoader for validation data.
+            test_dataloader (DataLoader | DictDataLoader |  None): DataLoader for testing data.
+        """
+        raise NotImplementedError("Writers must implement a log_model method.")
+
+    def print_metrics(self, result: OptimizeResult) -> None:
+        """Prints the metrics and loss in a readable format.
+
+        Args:
+            result (OptimizeResult): The optimization results to display.
+        """
+
+        # Find the key in result.metrics that contains "loss" (case-insensitive)
+        loss_key = next((k for k in result.metrics if "loss" in k.lower()), None)
+        if loss_key:
+            loss_value = result.metrics[loss_key]
+            msg = f"Iteration {result.iteration: >7} | {loss_key.title()}: {loss_value:.7f} -"
+        else:
+            msg = f"Iteration {result.iteration: >7} | Loss: None -"
+        msg += " ".join([f"{k}: {v:.7f}" for k, v in result.metrics.items() if k != loss_key])
+        print(msg)
+
+
+class TensorBoardWriter(BaseWriter):
+    """Writer for logging to TensorBoard.
+
+    Attributes:
+        writer (SummaryWriter): The TensorBoard SummaryWriter instance.
+    """
+
+    def __init__(self) -> None:
+        self.writer = None
+
+    def open(self, config: TrainConfig, iteration: int | None = None) -> SummaryWriter:
+        """
+        Opens the TensorBoard writer.
+
+        Args:
+            config: Configuration object containing settings for logging.
+            iteration (int, optional): The iteration step to start logging from.
+                Defaults to None.
+
+        Returns:
+            SummaryWriter: The initialized TensorBoard writer.
+        """
+        log_dir = str(config.log_folder)
+        purge_step = iteration if isinstance(iteration, int) else None
+        self.writer = SummaryWriter(log_dir=log_dir, purge_step=purge_step)
+        return self.writer
+
+    def close(self) -> None:
+        """Closes the TensorBoard writer."""
+        if self.writer:
+            self.writer.close()
+
+    def write(self, result: OptimizeResult) -> None:
+        """
+        Logs the results of the current iteration to TensorBoard.
+
+        Args:
+            result (OptimizeResult): The optimization results to log.
+        """
+        # Not writing loss as loss is available in the metrics
+        # if result.loss is not None:
+        #     self.writer.add_scalar("loss", float(result.loss), result.iteration)
+        if self.writer:
+            for key, value in result.metrics.items():
+                self.writer.add_scalar(key, value, result.iteration)
+        else:
+            raise RuntimeError(
+                "The writer is not initialized."
+                "Please call the 'writer.open()' method before writing"
+            )
+
+    def log_hyperparams(self, hyperparams: dict) -> None:
+        """
+        Logs hyperparameters to TensorBoard.
+
+        Args:
+            hyperparams (dict): A dictionary of hyperparameters to log.
+        """
+        if self.writer:
+            self.writer.add_hparams(hyperparams, {})
+        else:
+            raise RuntimeError(
+                "The writer is not initialized."
+                "Please call the 'writer.open()' method before writing"
+            )
+
+    def plot(
+        self,
+        model: Module,
+        iteration: int,
+        plotting_functions: tuple[PlottingFunction, ...],
+    ) -> None:
+        """
+        Logs plots of the model using provided plotting functions.
+
+        Args:
+            model (Module): The model to plot.
+            iteration (int): The current iteration number.
+            plotting_functions (tuple[PlottingFunction, ...]): Functions used
+                to generate plots.
+        """
+        if self.writer:
+            for pf in plotting_functions:
+                descr, fig = pf(model, iteration)
+                self.writer.add_figure(descr, fig, global_step=iteration)
+        else:
+            raise RuntimeError(
+                "The writer is not initialized."
+                "Please call the 'writer.open()' method before writing"
+            )
+
+    def log_model(
+        self,
+        model: Module,
+        train_dataloader: DataLoader | DictDataLoader | None = None,
+        val_dataloader: DataLoader | DictDataLoader | None = None,
+        test_dataloader: DataLoader | DictDataLoader | None = None,
+    ) -> None:
+        """
+        Logs the model.
+
+        Currently not supported by TensorBoard.
+
+        Args:
+            model (Module): The model to log.
+            train_dataloader (DataLoader | DictDataLoader |  None): DataLoader for training data.
+            val_dataloader (DataLoader | DictDataLoader |  None): DataLoader for validation data.
+            test_dataloader (DataLoader | DictDataLoader |  None): DataLoader for testing data.
+        """
+        logger.warning("Model logging is not supported by tensorboard. No model will be logged.")
+
+
+class MLFlowWriter(BaseWriter):
+    """
+    Writer for logging to MLflow.
+
+    Attributes:
+        run: The active MLflow run.
+        mlflow: The MLflow module.
+    """
+
+    def __init__(self) -> None:
+        try:
+            from mlflow.entities import Run
+        except ImportError:
+            raise ImportError(
+                "mlflow is not installed. Please install qadence with the mlflow feature: "
+                "`pip install qadence[mlflow]`."
+            )
+
+        self.run: Run
+        self.mlflow: ModuleType
+
+    def open(self, config: TrainConfig, iteration: int | None = None) -> ModuleType | None:
+        """
+        Opens the MLflow writer and initializes an MLflow run.
+
+        Args:
+            config: Configuration object containing settings for logging.
+            iteration (int, optional): The iteration step to start logging from.
+                Defaults to None.
+
+        Returns:
+            mlflow: The MLflow module instance.
+        """
+        import mlflow
+
+        self.mlflow = mlflow
+        tracking_uri = os.getenv("MLFLOW_TRACKING_URI", "")
+        experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME", str(uuid4()))
+        run_name = os.getenv("MLFLOW_RUN_NAME", str(uuid4()))
+
+        if self.mlflow:
+            self.mlflow.set_tracking_uri(tracking_uri)
+
+            # Create or get the experiment
+            exp_filter_string = f"name = '{experiment_name}'"
+            experiments = self.mlflow.search_experiments(filter_string=exp_filter_string)
+            if not experiments:
+                self.mlflow.create_experiment(name=experiment_name)
+
+            self.mlflow.set_experiment(experiment_name)
+            self.run = self.mlflow.start_run(run_name=run_name, nested=False)
+
+        return self.mlflow
+
+    def close(self) -> None:
+        """Closes the MLflow run."""
+        if self.run:
+            self.mlflow.end_run()
+
+    def write(self, result: OptimizeResult) -> None:
+        """
+        Logs the results of the current iteration to MLflow.
+
+        Args:
+            result (OptimizeResult): The optimization results to log.
+        """
+        # Not writing loss as loss is available in the metrics
+        # if result.loss is not None:
+        #     self.mlflow.log_metric("loss", float(result.loss), step=result.iteration)
+        if self.mlflow:
+            self.mlflow.log_metrics(result.metrics, step=result.iteration)
+        else:
+            raise RuntimeError(
+                "The writer is not initialized."
+                "Please call the 'writer.open()' method before writing"
+            )
+
+    def log_hyperparams(self, hyperparams: dict) -> None:
+        """
+        Logs hyperparameters to MLflow.
+
+        Args:
+            hyperparams (dict): A dictionary of hyperparameters to log.
+        """
+        if self.mlflow:
+            self.mlflow.log_params(hyperparams)
+        else:
+            raise RuntimeError(
+                "The writer is not initialized."
+                "Please call the 'writer.open()' method before writing"
+            )
+
+    def plot(
+        self,
+        model: Module,
+        iteration: int,
+        plotting_functions: tuple[PlottingFunction, ...],
+    ) -> None:
+        """
+        Logs plots of the model using provided plotting functions.
+
+        Args:
+            model (Module): The model to plot.
+            iteration (int): The current iteration number.
+            plotting_functions (tuple[PlottingFunction, ...]): Functions used
+                to generate plots.
+        """
+        if self.mlflow:
+            for pf in plotting_functions:
+                descr, fig = pf(model, iteration)
+                self.mlflow.log_figure(fig, descr)
+        else:
+            raise RuntimeError(
+                "The writer is not initialized."
+                "Please call the 'writer.open()' method before writing"
+            )
+
+    def get_signature_from_dataloader(
+        self, model: Module, dataloader: DataLoader | DictDataLoader | None
+    ) -> Any:
+        """
+        Infers the signature of the model based on the input data from the dataloader.
+
+        Args:
+            model (Module): The model to use for inference.
+            dataloader (DataLoader | DictDataLoader |  None): DataLoader for model inputs.
+
+        Returns:
+            Optional[Any]: The inferred signature, if available.
+        """
+        from mlflow.models import infer_signature
+
+        if dataloader is None:
+            return None
+
+        xs: InputData
+        xs, *_ = next(iter(dataloader))
+        preds = model(xs)
+
+        if isinstance(xs, Tensor):
+            xs = xs.detach().cpu().numpy()
+            preds = preds.detach().cpu().numpy()
+            return infer_signature(xs, preds)
+
+        return None
+
+    def log_model(
+        self,
+        model: Module,
+        train_dataloader: DataLoader | DictDataLoader | None = None,
+        val_dataloader: DataLoader | DictDataLoader | None = None,
+        test_dataloader: DataLoader | DictDataLoader | None = None,
+    ) -> None:
+        """
+        Logs the model and its signature to MLflow using the provided data loaders.
+
+        Args:
+            model (Module): The model to log.
+            train_dataloader (DataLoader | DictDataLoader |  None): DataLoader for training data.
+            val_dataloader (DataLoader | DictDataLoader |  None): DataLoader for validation data.
+            test_dataloader (DataLoader | DictDataLoader |  None): DataLoader for testing data.
+        """
+        if not self.mlflow:
+            raise RuntimeError(
+                "The writer is not initialized."
+                "Please call the 'writer.open()' method before writing"
+            )
+
+        signatures = self.get_signature_from_dataloader(model, train_dataloader)
+        self.mlflow.pytorch.log_model(model, artifact_path="model", signature=signatures)
+
+
+# Writer registry
+WRITER_REGISTRY = {
+    ExperimentTrackingTool.TENSORBOARD: TensorBoardWriter,
+    ExperimentTrackingTool.MLFLOW: MLFlowWriter,
+}
+
+
+def get_writer(tracking_tool: ExperimentTrackingTool) -> BaseWriter:
+    """Factory method to get the appropriate writer based on the tracking tool.
+
+    Args:
+        tracking_tool (ExperimentTrackingTool): The experiment tracking tool to use.
+
+    Returns:
+        BaseWriter: An instance of the appropriate writer.
+    """
+    writer_class = WRITER_REGISTRY.get(tracking_tool)
+    if writer_class:
+        return writer_class()
+    else:
+        raise ValueError(f"Unsupported tracking tool: {tracking_tool}")