qadence 1.7.3__py3-none-any.whl → 1.7.5__py3-none-any.whl

qadence/ml_tools/train_grad.py

@@ -3,16 +3,10 @@ from __future__ import annotations
 import importlib
 import math
 from logging import getLogger
-from typing import Callable, Union
-
-from rich.progress import (
-    BarColumn,
-    Progress,
-    TaskProgressColumn,
-    TextColumn,
-    TimeRemainingColumn,
-)
-from torch import complex128, float32, float64
+from typing import Any, Callable, Union
+
+from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeRemainingColumn
+from torch import Tensor, complex128, float32, float64
 from torch import device as torch_device
 from torch import dtype as torch_dtype
 from torch.nn import DataParallel, Module
@@ -20,8 +14,8 @@ from torch.optim import Optimizer
 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, data_to_device
+from qadence.ml_tools.config import Callback, TrainConfig
+from qadence.ml_tools.data import DictDataLoader, OptimizeResult, data_to_device
 from qadence.ml_tools.optimize_step import optimize_step
 from qadence.ml_tools.printing import (
     log_model_tracker,
@@ -166,107 +160,178 @@ def train(
 
     best_val_loss = math.inf
 
+    if not ((dataloader is None) or isinstance(dataloader, (DictDataLoader, DataLoader))):
+        raise NotImplementedError(
+            f"Unsupported dataloader type: {type(dataloader)}. "
+            "You can use e.g. `qadence.ml_tools.to_dataloader` to build a dataloader."
+        )
+
+    def next_loss_iter(dl_iter: Union[None, DataLoader, DictDataLoader]) -> Any:
+        """Get loss on the next batch of a dataloader.
+
+            loaded on device if not None.
+
+        Args:
+            dl_iter (Union[None, DataLoader, DictDataLoader]): Dataloader.
+
+        Returns:
+            Any: Loss value
+        """
+        xs = next(dl_iter) if dl_iter is not None else None
+        xs_to_device = data_to_device(xs, device=device, dtype=data_dtype)
+        return loss_fn(model, xs_to_device)
+
+    # populate callbacks with already available internal functions
+    # printing, writing and plotting
+    callbacks = config.callbacks
+
+    # printing
+    if config.verbose and config.print_every > 0:
+        # Note that the loss returned by optimize_step
+        # is the value before doing the training step
+        # which is printed accordingly by the previous iteration number
+        callbacks += [
+            Callback(
+                lambda opt_res: print_metrics(opt_res.loss, opt_res.metrics, opt_res.iteration - 1),
+                called_every=config.print_every,
+                call_after_opt=True,
+            )
+        ]
+
+    # plotting
+    callbacks += [
+        Callback(
+            lambda opt_res: plot_tracker(
+                writer,
+                opt_res.model,
+                opt_res.iteration,
+                config.plotting_functions,
+                tracking_tool=config.tracking_tool,
+            ),
+            called_every=config.plot_every,
+            call_before_opt=True,
+        )
+    ]
+
+    # writing metrics
+    callbacks += [
+        Callback(
+            lambda opt_res: write_tracker(
+                writer,
+                opt_res.loss,
+                opt_res.metrics,
+                opt_res.iteration,
+                tracking_tool=config.tracking_tool,
+            ),
+            called_every=config.write_every,
+            call_before_opt=False,
+            call_after_opt=True,
+            call_during_eval=True,
+        )
+    ]
+
+    # checkpointing
+    if config.folder and config.checkpoint_every > 0 and not config.checkpoint_best_only:
+        callbacks += [
+            Callback(
+                lambda opt_res: write_checkpoint(
+                    config.folder,  # type: ignore[arg-type]
+                    opt_res.model,
+                    opt_res.optimizer,
+                    opt_res.iteration,
+                ),
+                called_every=config.checkpoint_every,
+                call_before_opt=False,
+                call_after_opt=True,
+            )
+        ]
+
+    if config.folder and config.checkpoint_best_only:
+        callbacks += [
+            Callback(
+                lambda opt_res: write_checkpoint(
+                    config.folder,  # type: ignore[arg-type]
+                    opt_res.model,
+                    opt_res.optimizer,
+                    "best",
+                ),
+                called_every=config.checkpoint_every,
+                call_before_opt=True,
+                call_after_opt=True,
+                call_during_eval=True,
+            )
+        ]
+
+    def run_callbacks(callback_iterable: list[Callback], opt_res: OptimizeResult) -> None:
+        for callback in callback_iterable:
+            callback(opt_res)
+
+    callbacks_before_opt = [
+        callback
+        for callback in callbacks
+        if callback.call_before_opt and not callback.call_during_eval
+    ]
+    callbacks_before_opt_eval = [
+        callback for callback in callbacks if callback.call_before_opt and callback.call_during_eval
+    ]
+
     with progress:
         dl_iter = iter(dataloader) if dataloader is not None else None
 
         # Initial validation evaluation
         try:
+            opt_result = OptimizeResult(init_iter, model, optimizer)
             if perform_val:
                 dl_iter_val = iter(val_dataloader) if val_dataloader is not None else None
-                xs = next(dl_iter_val)
-                xs_to_device = data_to_device(xs, device=device, dtype=data_dtype)
-                best_val_loss, metrics = loss_fn(model, xs_to_device)
-
+                best_val_loss, metrics, *_ = next_loss_iter(dl_iter_val)
                 metrics["val_loss"] = best_val_loss
-                write_tracker(writer, None, metrics, init_iter, tracking_tool=config.tracking_tool)
-
-            if config.folder:
-                if config.checkpoint_best_only:
-                    write_checkpoint(config.folder, model, optimizer, iteration="best")
-                else:
-                    write_checkpoint(config.folder, model, optimizer, init_iter)
+                opt_result.metrics = metrics
+                run_callbacks(callbacks_before_opt_eval, opt_result)
 
-            plot_tracker(
-                writer,
-                model,
-                init_iter,
-                config.plotting_functions,
-                tracking_tool=config.tracking_tool,
-            )
+            run_callbacks(callbacks_before_opt, opt_result)
 
         except KeyboardInterrupt:
             logger.info("Terminating training gracefully after the current iteration.")
 
         # outer epoch loop
         init_iter += 1
+        callbacks_end_epoch = [
+            callback
+            for callback in callbacks
+            if callback.call_end_epoch and not callback.call_during_eval
+        ]
+        callbacks_end_epoch_eval = [
+            callback
+            for callback in callbacks
+            if callback.call_end_epoch and callback.call_during_eval
+        ]
         for iteration in progress.track(range(init_iter, init_iter + config.max_iter)):
             try:
                 # in case there is not data needed by the model
                 # this is the case, for example, of quantum models
                 # which do not have classical input data (e.g. chemistry)
-                if dataloader is None:
-                    loss, metrics = optimize_step(
-                        model=model,
-                        optimizer=optimizer,
-                        loss_fn=loss_fn,
-                        xs=None,
-                        device=device,
-                        dtype=data_dtype,
-                    )
+                loss, metrics = optimize_step(
+                    model=model,
+                    optimizer=optimizer,
+                    loss_fn=loss_fn,
+                    xs=None if dataloader is None else next(dl_iter),  # type: ignore[arg-type]
+                    device=device,
+                    dtype=data_dtype,
+                )
+                if isinstance(loss, Tensor):
                     loss = loss.item()
+                opt_result = OptimizeResult(iteration, model, optimizer, loss, metrics)
+                run_callbacks(callbacks_end_epoch, opt_result)
 
-                elif isinstance(dataloader, (DictDataLoader, DataLoader)):
-                    loss, metrics = optimize_step(
-                        model=model,
-                        optimizer=optimizer,
-                        loss_fn=loss_fn,
-                        xs=next(dl_iter),  # type: ignore[arg-type]
-                        device=device,
-                        dtype=data_dtype,
-                    )
-
-                else:
-                    raise NotImplementedError(
-                        f"Unsupported dataloader type: {type(dataloader)}. "
-                        "You can use e.g. `qadence.ml_tools.to_dataloader` to build a dataloader."
-                    )
-
-                if iteration % config.print_every == 0 and config.verbose:
-                    # Note that the loss returned by optimize_step
-                    # is the value before doing the training step
-                    # which is printed accordingly by the previous iteration number
-                    print_metrics(loss, metrics, iteration - 1)
-
-                if iteration % config.write_every == 0:
-                    write_tracker(
-                        writer, loss, metrics, iteration, tracking_tool=config.tracking_tool
-                    )
-
-                if iteration % config.plot_every == 0:
-                    plot_tracker(
-                        writer,
-                        model,
-                        iteration,
-                        config.plotting_functions,
-                        tracking_tool=config.tracking_tool,
-                    )
                 if perform_val:
                     if iteration % config.val_every == 0:
-                        xs = next(dl_iter_val)
-                        xs_to_device = data_to_device(xs, device=device, dtype=data_dtype)
-                        val_loss, *_ = loss_fn(model, xs_to_device)
+                        val_loss, *_ = next_loss_iter(dl_iter_val)
                         if config.validation_criterion(val_loss, best_val_loss, config.val_epsilon):  # type: ignore[misc]
                             best_val_loss = val_loss
-                            if config.folder and config.checkpoint_best_only:
-                                write_checkpoint(config.folder, model, optimizer, iteration="best")
                             metrics["val_loss"] = val_loss
-                            write_tracker(
-                                writer, loss, metrics, iteration, tracking_tool=config.tracking_tool
-                            )
+                            opt_result.metrics = metrics
 
-                if config.folder:
-                    if iteration % config.checkpoint_every == 0 and not config.checkpoint_best_only:
-                        write_checkpoint(config.folder, model, optimizer, iteration)
+                            run_callbacks(callbacks_end_epoch_eval, opt_result)
 
             except KeyboardInterrupt:
                 logger.info("Terminating training gracefully after the current iteration.")
@@ -275,21 +340,16 @@ def train(
         # Handling printing the last training loss
         # as optimize_step does not give the loss value at the last iteration
         try:
-            xs = next(dl_iter) if dataloader is not None else None  # type: ignore[arg-type]
-            xs_to_device = data_to_device(xs, device=device, dtype=data_dtype)
-            loss, metrics, *_ = loss_fn(model, xs_to_device)
-            if dataloader is None:
-                loss = loss.item()
+            loss, metrics, *_ = next_loss_iter(dl_iter)
             if iteration % config.print_every == 0 and config.verbose:
                 print_metrics(loss, metrics, iteration)
 
         except KeyboardInterrupt:
             logger.info("Terminating training gracefully after the current iteration.")
 
-    # Final checkpointing and writing
-    if config.folder and not config.checkpoint_best_only:
-        write_checkpoint(config.folder, model, optimizer, iteration)
-    write_tracker(writer, loss, metrics, iteration, tracking_tool=config.tracking_tool)
+    # Final callbacks, by default checkpointing and writing
+    callbacks_after_opt = [callback for callback in callbacks if callback.call_after_opt]
+    run_callbacks(callbacks_after_opt, opt_result)
 
     # writing hyperparameters
     if config.hyperparams:

qadence/ml_tools/train_no_grad.py

@@ -6,20 +6,14 @@ from typing import Callable
 
 import nevergrad as ng
 from nevergrad.optimization.base import Optimizer as NGOptimizer
-from rich.progress import (
-    BarColumn,
-    Progress,
-    TaskProgressColumn,
-    TextColumn,
-    TimeRemainingColumn,
-)
+from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeRemainingColumn
 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
+from qadence.ml_tools.config import Callback, TrainConfig
+from qadence.ml_tools.data import DictDataLoader, OptimizeResult
 from qadence.ml_tools.parameters import get_parameters, set_parameters
 from qadence.ml_tools.printing import (
     log_model_tracker,
@@ -92,6 +86,12 @@ def train(
     params = get_parameters(model).detach().numpy()
     ng_params = ng.p.Array(init=params)
 
+    if not ((dataloader is None) or isinstance(dataloader, (DictDataLoader, DataLoader))):
+        raise NotImplementedError(
+            f"Unsupported dataloader type: {type(dataloader)}. "
+            "You can use e.g. `qadence.ml_tools.to_dataloader` to build a dataloader."
+        )
+
     # serial training
     # TODO: Add a parallelization using the num_workers argument in Nevergrad
     progress = Progress(
@@ -100,38 +100,85 @@ def train(
         TaskProgressColumn(),
         TimeRemainingColumn(elapsed_when_finished=True),
     )
-    with progress:
-        dl_iter = iter(dataloader) if dataloader is not None else None
-
-        for iteration in progress.track(range(init_iter, init_iter + config.max_iter)):
-            if dataloader is None:
-                loss, metrics, ng_params = _update_parameters(None, ng_params)
-
-            elif isinstance(dataloader, (DictDataLoader, DataLoader)):
-                data = next(dl_iter)  # type: ignore[arg-type]
-                loss, metrics, ng_params = _update_parameters(data, ng_params)
-
-            else:
-                raise NotImplementedError("Unsupported dataloader type!")
-
-            if iteration % config.print_every == 0 and config.verbose:
-                print_metrics(loss, metrics, iteration)
 
-            if iteration % config.write_every == 0:
-                write_tracker(writer, loss, metrics, iteration, tracking_tool=config.tracking_tool)
-
-            if iteration % config.plot_every == 0:
-                plot_tracker(
+    # populate callbacks with already available internal functions
+    # printing, writing and plotting
+    callbacks = config.callbacks
+
+    # printing
+    if config.verbose and config.print_every > 0:
+        callbacks += [
+            Callback(
+                lambda opt_res: print_metrics(opt_res.loss, opt_res.metrics, opt_res.iteration),
+                called_every=config.print_every,
+            )
+        ]
+
+    # writing metrics
+    if config.write_every > 0:
+        callbacks += [
+            Callback(
+                lambda opt_res: write_tracker(
+                    writer,
+                    opt_res.loss,
+                    opt_res.metrics,
+                    opt_res.iteration,
+                    tracking_tool=config.tracking_tool,
+                ),
+                called_every=config.write_every,
+                call_after_opt=True,
+            )
+        ]
+
+    # plot tracker
+    if config.plot_every > 0:
+        callbacks += [
+            Callback(
+                lambda opt_res: plot_tracker(
                     writer,
-                    model,
-                    iteration,
+                    opt_res.model,
+                    opt_res.iteration,
                     config.plotting_functions,
                     tracking_tool=config.tracking_tool,
-                )
+                ),
+                called_every=config.plot_every,
+            )
+        ]
+
+    # checkpointing
+    if config.folder and config.checkpoint_every > 0:
+        callbacks += [
+            Callback(
+                lambda opt_res: write_checkpoint(
+                    config.folder,  # type: ignore[arg-type]
+                    opt_res.model,
+                    opt_res.optimizer,
+                    opt_res.iteration,
+                ),
+                called_every=config.checkpoint_every,
+                call_after_opt=True,
+            )
+        ]
+
+    def run_callbacks(callback_iterable: list[Callback], opt_res: OptimizeResult) -> None:
+        for callback in callback_iterable:
+            callback(opt_res)
+
+    callbacks_end_opt = [
+        callback
+        for callback in callbacks
+        if callback.call_end_epoch and not callback.call_during_eval
+    ]
+
+    with progress:
+        dl_iter = iter(dataloader) if dataloader is not None else None
 
-            if config.folder:
-                if iteration % config.checkpoint_every == 0:
-                    write_checkpoint(config.folder, model, optimizer, iteration)
+        for iteration in progress.track(range(init_iter, init_iter + config.max_iter)):
+            loss, metrics, ng_params = _update_parameters(
+                None if dataloader is None else next(dl_iter), ng_params  # type: ignore[arg-type]
+            )
+            opt_result = OptimizeResult(iteration, model, optimizer, loss, metrics)
+            run_callbacks(callbacks_end_opt, opt_result)
 
             if iteration >= init_iter + config.max_iter:
                 break
@@ -143,10 +190,9 @@ def train(
     if config.log_model:
         log_model_tracker(writer, model, dataloader, tracking_tool=config.tracking_tool)
 
-    # Final writing and checkpointing
-    if config.folder:
-        write_checkpoint(config.folder, model, optimizer, iteration)
-    write_tracker(writer, loss, metrics, iteration, tracking_tool=config.tracking_tool)
+    # Final callbacks
+    callbacks_after_opt = [callback for callback in callbacks if callback.call_after_opt]
+    run_callbacks(callbacks_after_opt, opt_result)
 
     # close tracker
     if config.tracking_tool == ExperimentTrackingTool.TENSORBOARD:

qadence/model.py

@@ -353,11 +353,11 @@ class QuantumModel(nn.Module):
         """
         raise NotImplementedError("The overlap method is not implemented for this model.")
 
-    def _to_dict(self, save_params: bool = False) -> dict[str, Any]:
+    def _to_dict(self, save_params: bool = True) -> dict[str, Any]:
         """Convert QuantumModel to a dictionary for serialization.
 
         Arguments:
-            save_params: Optionally save parameters. Defaults to False.
+            save_params: Save parameters. Defaults to True.
 
         Returns:
             The dictionary
@@ -382,7 +382,7 @@ class QuantumModel(nn.Module):
             }
             param_dict_conv = {}
             if save_params:
-                param_dict_conv = {name: param.data for name, param in self._params.items()}
+                param_dict_conv = {name: param for name, param in self._params.items()}
             d = {self.__class__.__name__: d, "param_dict": param_dict_conv}
             logger.debug(f"{self.__class__.__name__} serialized to {d}.")
         except Exception as e:
@@ -432,6 +432,50 @@ class QuantumModel(nn.Module):
 
         return qm
 
+    def load_params_from_dict(self, d: dict, strict: bool = True) -> None:
+        """Copy parameters from dictionary into this QuantumModel.
+
+        Unlike :meth:`~qadence.QuantumModel.from_dict`, this method does not create a new
+        QuantumModel instance, but rather loads the parameters into the same QuantumModel.
+        The behaviour of this method is similar to :meth:`~torch.nn.Module.load_state_dict`.
+
+        The dictionary is assumed to have the format as saved via
+        :meth:`~qadence.QuantumModel.to_dict`
+
+        Args:
+            d (dict): The dictionary
+            strict (bool, optional):
+                Whether to strictly enforce that the parameter keys in the dictionary and
+                in the model match exactly. Default: ``True``.
+        """
+        param_dict = d["param_dict"]
+        missing_keys = set(self._params.keys()) - set(param_dict.keys())
+        unexpected_keys = set(param_dict.keys()) - set(self._params.keys())
+
+        if strict:
+            error_msgs = []
+            if len(unexpected_keys) > 0:
+                error_msgs.append(f"Unexpected key(s) in dictionary: {unexpected_keys}")
+            if len(missing_keys) > 0:
+                error_msgs.append(f"Missing key(s) in dictionary: {missing_keys}")
+            if len(error_msgs) > 0:
+                errors_string = "\n\t".join(error_msgs)
+                raise RuntimeError(
+                    f"Error(s) loading the parameter dictionary due to: \n\t{errors_string}\n"
+                    "This error was thrown because the `strict` argument is set `True`."
+                    "If you don't need the parameter keys of the dictionary to exactly match "
+                    "the model parameters, set `strict=False`."
+                )
+
+        for n, param in param_dict.items():
+            try:
+                with torch.no_grad():
+                    self._params[n].copy_(
+                        torch.nn.Parameter(param, requires_grad=param.requires_grad)
+                    )
+            except Exception as e:
+                logger.warning(f"Unable to load parameter {n} from dictionary due to {e}.")
+
     def save(
         self, folder: str | Path, file_name: str = "quantum_model.pt", save_params: bool = True
     ) -> None:

qadence/types.py

@@ -156,9 +156,9 @@ class ReuploadScaling(StrEnum):
 class MultivariateStrategy(StrEnum):
     """Multivariate strategy for feature maps."""
 
-    PARALLEL = "parallel"
+    PARALLEL = "Parallel"
     """Parallel strategy."""
-    SERIES = "SERIES"
+    SERIES = "Series"
     """Serial strategy."""
 
 

{qadence-1.7.3.dist-info → qadence-1.7.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: qadence
-Version: 1.7.3
+Version: 1.7.5
 Summary: Pasqal interface for circuit-based quantum computing SDKs
 Author-email: Aleksander Wennersteen <aleksander.wennersteen@pasqal.com>, Gert-Jan Both <gert-jan.both@pasqal.com>, Niklas Heim <niklas.heim@pasqal.com>, Mario Dagrada <mario.dagrada@pasqal.com>, Vincent Elfving <vincent.elfving@pasqal.com>, Dominik Seitz <dominik.seitz@pasqal.com>, Roland Guichard <roland.guichard@pasqal.com>, "Joao P. Moutinho" <joao.moutinho@pasqal.com>, Vytautas Abramavicius <vytautas.abramavicius@pasqal.com>, Gergana Velikova <gergana.velikova@pasqal.com>, Eduardo Maschio <eduardo.maschio@pasqal.com>, Smit Chaudhary <smit.chaudhary@pasqal.com>, Ignacio Fernández Graña <ignacio.fernandez-grana@pasqal.com>, Charles Moussa <charles.moussa@pasqal.com>, Giorgio Tosti Balducci <giorgio.tosti-balducci@pasqal.com>
 License: Apache 2.0
@@ -22,7 +22,7 @@ Requires-Dist: matplotlib
 Requires-Dist: nevergrad
 Requires-Dist: numpy
 Requires-Dist: openfermion
-Requires-Dist: pyqtorch==1.3.2
+Requires-Dist: pyqtorch==1.4.3
 Requires-Dist: pyyaml
 Requires-Dist: rich
 Requires-Dist: scipy
@@ -45,7 +45,7 @@ Requires-Dist: nvidia-pyindex; extra == 'dlprof'
 Provides-Extra: horqrux
 Requires-Dist: einops; extra == 'horqrux'
 Requires-Dist: flax; extra == 'horqrux'
-Requires-Dist: horqrux==0.6.0; extra == 'horqrux'
+Requires-Dist: horqrux==0.6.2; extra == 'horqrux'
 Requires-Dist: jax; extra == 'horqrux'
 Requires-Dist: jaxopt; extra == 'horqrux'
 Requires-Dist: optax; extra == 'horqrux'
@@ -57,7 +57,7 @@ Requires-Dist: mlflow; extra == 'mlflow'
 Provides-Extra: protocols
 Requires-Dist: qadence-protocols; extra == 'protocols'
 Provides-Extra: pulser
-Requires-Dist: pasqal-cloud==0.11.1; extra == 'pulser'
+Requires-Dist: pasqal-cloud==0.11.3; extra == 'pulser'
 Requires-Dist: pulser-core==0.19.0; extra == 'pulser'
 Requires-Dist: pulser-simulation==0.19.0; extra == 'pulser'
 Provides-Extra: visualization