QML tools

ML Tools

This module implements a Trainer class for torch Modules and QuantumModel. It also implements the QNN class and callbacks that can be used with the trainer module.

Trainer(model, optimizer, config, loss_fn='mse', train_dataloader=None, val_dataloader=None, test_dataloader=None, optimize_step=optimize_step, device=None, dtype=None, max_batches=None)

Bases: BaseTrainer

Trainer class to manage and execute training, validation, and testing loops for a model (eg.

QNN).

This class handles the overall training process, including: - Managing epochs and steps - Handling data loading and batching - Computing and updating gradients - Logging and monitoring training metrics

ATTRIBUTE	DESCRIPTION
current_epoch	The current epoch number. TYPE: int
global_step	The global step across all epochs. TYPE: int
log_device	Device for logging, default is "cpu". TYPE: str
device	Device used for computation. TYPE: device
dtype	Data type used for computation. TYPE: dtype | None
data_dtype	Data type for data. Depends on the model's data type. TYPE: dtype | None

Inherited Attributes

use_grad (bool): Indicates if gradients are used for optimization. Default is True.

model (nn.Module): The neural network model. optimizer (optim.Optimizer | NGOptimizer | None): The optimizer for training. config (TrainConfig): The configuration settings for training. 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.

optimize_step (Callable): Function for performing an optimization step. loss_fn (Callable): loss function to use.

num_training_batches (int): Number of training batches. num_validation_batches (int): Number of validation batches. num_test_batches (int): Number of test batches.

state (str): Current state in the training process

Default training routine

for epoch in max_iter + 1:
    # Training
    for batch in train_batches:
        train model
    # Validation
    if val_every % epoch == 0:
        for batch in val_batches:
            train model

Notes

• In case of InfiniteTensorDataset, number of batches = 1.
• In case of TensorDataset, number of batches are default.
• Training is run for max_iter + 1 epochs. Epoch 0 logs untrained model.
• Please look at the CallbackManager initialize_callbacks method to review the default logging behavior.

Examples:

import torch
from torch.optim import SGD
from qadence import (
    feature_map,
    hamiltonian_factory,
    hea,
    QNN,
    QuantumCircuit,
    TrainConfig,
    Z,
)
from qadence.ml_tools.trainer import Trainer
from qadence.ml_tools.optimize_step import optimize_step
from qadence.ml_tools import TrainConfig
from qadence.ml_tools.data import to_dataloader

# Initialize the model
n_qubits = 2
fm = feature_map(n_qubits)
ansatz = hea(n_qubits=n_qubits, depth=2)
observable = hamiltonian_factory(n_qubits, detuning=Z)
circuit = QuantumCircuit(n_qubits, fm, ansatz)
model = QNN(circuit, observable, backend="pyqtorch", diff_mode="ad")

# Set up the optimizer
optimizer = SGD(model.parameters(), lr=0.001)

# Use TrainConfig for configuring the training process
config = TrainConfig(
    max_iter=100,
    print_every=10,
    write_every=10,
    checkpoint_every=10,
    val_every=10
)

# Create the Trainer instance with TrainConfig
trainer = Trainer(
    model=model,
    optimizer=optimizer,
    config=config,
    loss_fn="mse",
    optimize_step=optimize_step
)

batch_size = 25
x = torch.linspace(0, 1, 32).reshape(-1, 1)
y = torch.sin(x)
train_loader = to_dataloader(x, y, batch_size=batch_size, infinite=True)
val_loader = to_dataloader(x, y, batch_size=batch_size, infinite=False)

# Train the model
model, optimizer = trainer.fit(train_loader, val_loader)

This also supports both gradient based and gradient free optimization. The default support is for gradient based optimization.

Notes:

• set_use_grad() (class level):This method is used to set the global use_grad flag, controlling whether the trainer uses gradient-based optimization.

  # gradient based
  Trainer.set_use_grad(True)
  
  # gradient free
  Trainer.set_use_grad(False)
  

• Context Managers (instance level): enable_grad_opt() and disable_grad_opt() are context managers that temporarily switch the optimization mode for specific code blocks. This is useful when you want to mix gradient-based and gradient-free optimization in the same training process.

  # gradient based
  with trainer.enable_grad_opt(optimizer):
      trainer.fit()
  
  # gradient free
  with trainer.disable_grad_opt(ng_optimizer):
      trainer.fit()
  

Examples

Gradient based optimization example Usage:

from torch import optim
optimizer = optim.SGD(model.parameters(), lr=0.01)

Trainer.set_use_grad(True)
trainer = Trainer(
    model=model,
    optimizer=optimizer,
    config=config,
    loss_fn="mse"
)
trainer.fit(train_loader, val_loader)

or

trainer = Trainer(
    model=model,
    config=config,
    loss_fn="mse"
)
with trainer.enable_grad_opt(optimizer):
    trainer.fit(train_loader, val_loader)

Gradient free optimization example Usage:

import nevergrad as ng
from qadence.ml_tools.parameters import num_parameters
ng_optimizer = ng.optimizers.NGOpt(
                budget=config.max_iter, parametrization= num_parameters(model)
                )

Trainer.set_use_grad(False)
trainer = Trainer(
    model=model,
    optimizer=ng_optimizer,
    config=config,
    loss_fn="mse"
)
trainer.fit(train_loader, val_loader)

or

import nevergrad as ng
from qadence.ml_tools.parameters import num_parameters
ng_optimizer = ng.optimizers.NGOpt(
        budget=config.max_iter, parametrization= num_parameters(model)
        )

trainer = Trainer(
    model=model,
    config=config,
    loss_fn="mse"
)
with trainer.disable_grad_opt(ng_optimizer):
    trainer.fit(train_loader, val_loader)

Initializes the Trainer class.

PARAMETER	DESCRIPTION
model	The PyTorch model to train. TYPE: Module
optimizer	The optimizer for training. TYPE: Optimizer | Optimizer | None
config	Training configuration object. TYPE: TrainConfig
loss_fn	Loss function used for training. If not specified, default mse loss will be used. TYPE: str | Callable DEFAULT: 'mse'
train_dataloader	DataLoader for training data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
val_dataloader	DataLoader for validation data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
test_dataloader	DataLoader for test data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
optimize_step	Function to execute an optimization step. TYPE: Callable DEFAULT: optimize_step
device	Device to use for computation. TYPE: device DEFAULT: None
dtype	Data type for computation. TYPE: dtype DEFAULT: None
max_batches	Maximum number of batches to process per epoch. This is only valid in case of finite TensorDataset dataloaders. if max_batches is not None, the maximum number of batches used will be min(max_batches, len(dataloader.dataset)) In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: int | None DEFAULT: None

Source code in qadence/ml_tools/trainer.py

def __init__(
    self,
    model: nn.Module,
    optimizer: optim.Optimizer | NGOptimizer | None,
    config: TrainConfig,
    loss_fn: str | Callable = "mse",
    train_dataloader: DataLoader | DictDataLoader | None = None,
    val_dataloader: DataLoader | DictDataLoader | None = None,
    test_dataloader: DataLoader | DictDataLoader | None = None,
    optimize_step: Callable = optimize_step,
    device: torch_device | None = None,
    dtype: torch_dtype | None = None,
    max_batches: int | None = None,
):
    """
    Initializes the Trainer class.

    Args:
        model (nn.Module): The PyTorch model to train.
        optimizer (optim.Optimizer | NGOptimizer | None): The optimizer for training.
        config (TrainConfig): Training configuration object.
        loss_fn (str | Callable ): Loss function used for training.
            If not specified, default mse loss will be used.
        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 test data.
        optimize_step (Callable): Function to execute an optimization step.
        device (torch_device): Device to use for computation.
        dtype (torch_dtype): Data type for computation.
        max_batches (int | None): Maximum number of batches to process per epoch.
            This is only valid in case of finite TensorDataset dataloaders.
            if max_batches is not None, the maximum number of batches used will
            be min(max_batches, len(dataloader.dataset))
            In case of InfiniteTensorDataset only 1 batch per epoch is used.
    """
    super().__init__(
        model=model,
        optimizer=optimizer,
        config=config,
        loss_fn=loss_fn,
        optimize_step=optimize_step,
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        test_dataloader=test_dataloader,
        max_batches=max_batches,
    )
    self.current_epoch: int = 0
    self.global_step: int = 0
    self.log_device: str = "cpu" if device is None else device
    self.device: torch_device | None = device
    self.dtype: torch_dtype | None = dtype
    self.data_dtype: torch_dtype | None = None
    if self.dtype:
        self.data_dtype = float64 if (self.dtype == complex128) else float32

build_optimize_result(result)

Builds and stores the optimization result by calculating the average loss and metrics.

Result (or loss_metrics) can have multiple formats: - None Indicates no loss or metrics data is provided. - tuple[torch.Tensor, dict[str, Any]] A single tuple containing the loss tensor and metrics dictionary - at the end of batch. - list[tuple[torch.Tensor, dict[str, Any]]] A list of tuples for multiple batches. - list[list[tuple[torch.Tensor, dict[str, Any]]]] A list of lists of tuples, where each inner list represents metrics across multiple batches within an epoch.

PARAMETER	DESCRIPTION
result	(None | tuple[torch.Tensor, dict[Any, Any]] | list[tuple[torch.Tensor, dict[Any, Any]]] | list[list[tuple[torch.Tensor, dict[Any, Any]]]]) The loss and metrics data, which can have multiple formats TYPE: None | tuple[Tensor, dict[Any, Any]] | list[tuple[Tensor, dict[Any, Any]]] | list[list[tuple[Tensor, dict[Any, Any]]]]

RETURNS	DESCRIPTION
None	This method does not return anything. It sets self.opt_result with TYPE: None
None	the computed average loss and metrics.

Source code in qadence/ml_tools/trainer.py

def build_optimize_result(
    self,
    result: None
    | tuple[torch.Tensor, dict[Any, Any]]
    | list[tuple[torch.Tensor, dict[Any, Any]]]
    | list[list[tuple[torch.Tensor, dict[Any, Any]]]],
) -> None:
    """
    Builds and stores the optimization result by calculating the average loss and metrics.

    Result (or loss_metrics) can have multiple formats:
    - `None` Indicates no loss or metrics data is provided.
    - `tuple[torch.Tensor, dict[str, Any]]` A single tuple containing the loss tensor
        and metrics dictionary - at the end of batch.
    - `list[tuple[torch.Tensor, dict[str, Any]]]` A list of tuples for
        multiple batches.
    - `list[list[tuple[torch.Tensor, dict[str, Any]]]]` A list of lists of tuples,
    where each inner list represents metrics across multiple batches within an epoch.

    Args:
        result: (None |
                tuple[torch.Tensor, dict[Any, Any]] |
                list[tuple[torch.Tensor, dict[Any, Any]]] |
                list[list[tuple[torch.Tensor, dict[Any, Any]]]])
                    The loss and metrics data, which can have multiple formats

    Returns:
        None: This method does not return anything. It sets `self.opt_result` with
        the computed average loss and metrics.
    """
    loss_metrics = result
    if loss_metrics is None:
        loss = None
        metrics: dict[Any, Any] = {}
    elif isinstance(loss_metrics, tuple):
        # Single tuple case
        loss, metrics = loss_metrics
    else:
        last_epoch: list[tuple[torch.Tensor, dict[Any, Any]]] = []
        if isinstance(loss_metrics, list):
            # Check if it's a list of tuples
            if all(isinstance(item, tuple) for item in loss_metrics):
                last_epoch = cast(list[tuple[torch.Tensor, dict[Any, Any]]], loss_metrics)
            # Check if it's a list of lists of tuples
            elif all(isinstance(item, list) for item in loss_metrics):
                last_epoch = cast(
                    list[tuple[torch.Tensor, dict[Any, Any]]],
                    loss_metrics[-1] if loss_metrics else [],
                )
            else:
                raise ValueError(
                    "Invalid format for result: Expected None, tuple, list of tuples,"
                    " or list of lists of tuples."
                )

        if not last_epoch:
            loss, metrics = None, {}
        else:
            # Compute the average loss over the batches
            loss_tensor = torch.stack([loss_batch for loss_batch, _ in last_epoch])
            avg_loss = loss_tensor.mean()

            # Collect and average metrics for all batches
            metric_keys = last_epoch[0][1].keys()
            metrics_stacked: dict = {key: [] for key in metric_keys}

            for _, metrics_batch in last_epoch:
                for key in metric_keys:
                    value = metrics_batch[key]
                    metrics_stacked[key].append(value)

            avg_metrics = {key: torch.stack(metrics_stacked[key]).mean() for key in metric_keys}

            loss, metrics = avg_loss, avg_metrics

    # Store the optimization result
    self.opt_result = OptimizeResult(
        self.current_epoch, self.model_old, self.optimizer_old, loss, metrics
    )

fit(train_dataloader=None, val_dataloader=None)

Fits the model using the specified training configuration.

The dataloaders can be provided to train on new datasets, or the default dataloaders provided in the trainer will be used.

PARAMETER	DESCRIPTION
train_dataloader	DataLoader for training data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
val_dataloader	DataLoader for validation data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None

RETURNS	DESCRIPTION
tuple[Module, Optimizer]	tuple[nn.Module, optim.Optimizer]: The trained model and optimizer.

Source code in qadence/ml_tools/trainer.py

def fit(
    self,
    train_dataloader: DataLoader | DictDataLoader | None = None,
    val_dataloader: DataLoader | DictDataLoader | None = None,
) -> tuple[nn.Module, optim.Optimizer]:
    """
    Fits the model using the specified training configuration.

    The dataloaders can be provided to train on new datasets, or the default dataloaders
    provided in the trainer will be used.

    Args:
        train_dataloader (DataLoader | DictDataLoader |  None): DataLoader for training data.
        val_dataloader (DataLoader | DictDataLoader |  None): DataLoader for validation data.

    Returns:
        tuple[nn.Module, optim.Optimizer]: The trained model and optimizer.
    """
    if train_dataloader is not None:
        self.train_dataloader = train_dataloader
    if val_dataloader is not None:
        self.val_dataloader = val_dataloader

    self._fit_setup()
    self._train()
    self._fit_end()
    self.training_stage = TrainingStage("idle")
    return self.model, self.optimizer

run_test_batch(batch)

Runs a single test batch.

PARAMETER	DESCRIPTION
batch	Batch of data from the DataLoader. TYPE: tuple[Tensor, ...]

RETURNS	DESCRIPTION
tuple[Tensor, dict[str, Any]]	tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("test_batch")
def run_test_batch(
    self, batch: tuple[torch.Tensor, ...]
) -> tuple[torch.Tensor, dict[str, Any]]:
    """
    Runs a single test batch.

    Args:
        batch (tuple[torch.Tensor, ...]): Batch of data from the DataLoader.

    Returns:
        tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.
    """
    with torch.no_grad():
        loss_metrics = self.loss_fn(self.model, batch)
    return self._modify_batch_end_loss_metrics(loss_metrics)

run_train_batch(batch)

Runs a single training batch, performing optimization.

We use the step function to optimize the model based on use_grad. use_grad = True entails gradient based optimization, for which we use optimize_step function. use_grad = False entails gradient free optimization, for which we use update_ng_parameters function.

PARAMETER	DESCRIPTION
batch	Batch of data from the DataLoader. TYPE: tuple[Tensor, ...]

RETURNS	DESCRIPTION
tuple[Tensor, dict[str, Any]]	tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch. tuple of (loss, metrics)

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("train_batch")
def run_train_batch(
    self, batch: tuple[torch.Tensor, ...]
) -> tuple[torch.Tensor, dict[str, Any]]:
    """
    Runs a single training batch, performing optimization.

    We use the step function to optimize the model based on use_grad.
        use_grad = True entails gradient based optimization, for which we use
        optimize_step function.
        use_grad = False entails gradient free optimization, for which we use
        update_ng_parameters function.

    Args:
        batch (tuple[torch.Tensor, ...]): Batch of data from the DataLoader.

    Returns:
        tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.
            tuple of (loss, metrics)
    """

    if self.use_grad:
        # Perform gradient-based optimization
        loss_metrics = self.optimize_step(
            model=self.model,
            optimizer=self.optimizer,
            loss_fn=self.loss_fn,
            xs=batch,
            device=self.device,
            dtype=self.data_dtype,
        )
    else:
        # Perform optimization using Nevergrad
        loss, metrics, ng_params = update_ng_parameters(
            model=self.model,
            optimizer=self.optimizer,
            loss_fn=self.loss_fn,
            data=batch,
            ng_params=self.ng_params,  # type: ignore[arg-type]
        )
        self.ng_params = ng_params
        loss_metrics = loss, metrics

    return self._modify_batch_end_loss_metrics(loss_metrics)

run_training(dataloader)

Runs the training for a single epoch, iterating over multiple batches.

PARAMETER	DESCRIPTION
dataloader	DataLoader for training data. TYPE: DataLoader

RETURNS	DESCRIPTION
list[tuple[Tensor, dict[str, Any]]]	list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch. list -> tuples Training Batches -> (loss, metrics)

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("train_epoch")
def run_training(self, dataloader: DataLoader) -> list[tuple[torch.Tensor, dict[str, Any]]]:
    """
    Runs the training for a single epoch, iterating over multiple batches.

    Args:
        dataloader (DataLoader): DataLoader for training data.

    Returns:
        list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch.
            list                  -> tuples
            Training Batches      -> (loss, metrics)
    """
    self.model.train()
    train_epoch_loss_metrics = []
    # Quick Fix for iteration 0
    self._reset_model_and_opt()

    for batch in self._batch_iter(dataloader, self.num_training_batches):
        self.on_train_batch_start(batch)
        train_batch_loss_metrics = self.run_train_batch(batch)
        train_epoch_loss_metrics.append(train_batch_loss_metrics)
        self.on_train_batch_end(train_batch_loss_metrics)

    return train_epoch_loss_metrics

run_val_batch(batch)

Runs a single validation batch.

PARAMETER	DESCRIPTION
batch	Batch of data from the DataLoader. TYPE: tuple[Tensor, ...]

RETURNS	DESCRIPTION
tuple[Tensor, dict[str, Any]]	tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("val_batch")
def run_val_batch(self, batch: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, dict[str, Any]]:
    """
    Runs a single validation batch.

    Args:
        batch (tuple[torch.Tensor, ...]): Batch of data from the DataLoader.

    Returns:
        tuple[torch.Tensor, dict[str, Any]]: Loss and metrics for the batch.
    """
    with torch.no_grad():
        loss_metrics = self.loss_fn(self.model, batch)
    return self._modify_batch_end_loss_metrics(loss_metrics)

run_validation(dataloader)

Runs the validation loop for a single epoch, iterating over multiple batches.

PARAMETER	DESCRIPTION
dataloader	DataLoader for validation data. TYPE: DataLoader

RETURNS	DESCRIPTION
list[tuple[Tensor, dict[str, Any]]]	list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch. list -> tuples Validation Batches -> (loss, metrics)

Source code in qadence/ml_tools/trainer.py

@BaseTrainer.callback("val_epoch")
def run_validation(self, dataloader: DataLoader) -> list[tuple[torch.Tensor, dict[str, Any]]]:
    """
    Runs the validation loop for a single epoch, iterating over multiple batches.

    Args:
        dataloader (DataLoader): DataLoader for validation data.

    Returns:
        list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch.
            list                  -> tuples
            Validation Batches      -> (loss, metrics)
    """
    self.model.eval()
    val_epoch_loss_metrics = []

    for batch in self._batch_iter(dataloader, self.num_validation_batches):
        self.on_val_batch_start(batch)
        val_batch_loss_metrics = self.run_val_batch(batch)
        val_epoch_loss_metrics.append(val_batch_loss_metrics)
        self.on_val_batch_end(val_batch_loss_metrics)

    return val_epoch_loss_metrics

test(test_dataloader=None)

Runs the testing loop if a test DataLoader is provided.

if the test_dataloader is not provided, default test_dataloader defined in the Trainer class is used.

PARAMETER	DESCRIPTION
test_dataloader	DataLoader for test data. TYPE: DataLoader DEFAULT: None

RETURNS	DESCRIPTION
list[tuple[Tensor, dict[str, Any]]]	list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch. list -> tuples Test Batches -> (loss, metrics)

Source code in qadence/ml_tools/trainer.py

def test(self, test_dataloader: DataLoader = None) -> list[tuple[torch.Tensor, dict[str, Any]]]:
    """
    Runs the testing loop if a test DataLoader is provided.

    if the test_dataloader is not provided, default test_dataloader defined
    in the Trainer class is used.

    Args:
        test_dataloader (DataLoader): DataLoader for test data.

    Returns:
        list[tuple[torch.Tensor, dict[str, Any]]]: Loss and metrics for each batch.
            list                    -> tuples
            Test Batches            -> (loss, metrics)
    """
    if test_dataloader is not None:
        self.test_dataloader = test_dataloader

    self.model.eval()
    test_loss_metrics = []

    for batch in self._batch_iter(test_dataloader, self.num_training_batches):
        self.on_test_batch_start(batch)
        loss_metrics = self.run_test_batch(batch)
        test_loss_metrics.append(loss_metrics)
        self.on_test_batch_end(loss_metrics)

    return test_loss_metrics

AnsatzConfig(depth=1, ansatz_type=AnsatzType.HEA, ansatz_strategy=Strategy.DIGITAL, strategy_args=dict(), param_prefix='theta', tag=None) dataclass

ansatz_strategy: Strategy = Strategy.DIGITAL class-attribute instance-attribute

Ansatz strategy.

Strategy.DIGITAL for fully digital ansatz. Required if ansatz_type is AnsatzType.IIA. Strategy.SDAQC for analog entangling block. Strategy.RYDBERG for fully rydberg hea ansatz.

ansatz_type: AnsatzType = AnsatzType.HEA class-attribute instance-attribute

What type of ansatz.

AnsatzType.HEA for Hardware Efficient Ansatz. AnsatzType.IIA for Identity intialized Ansatz.

depth: int = 1 class-attribute instance-attribute

Number of layers of the ansatz.

param_prefix: str = 'theta' class-attribute instance-attribute

The base bame of the variational parameter.

strategy_args: dict = field(default_factory=dict) class-attribute instance-attribute

A dictionary containing keyword arguments to the function creating the ansatz.

Details about each below.

For Strategy.DIGITAL strategy, accepts the following: periodic (bool): if the qubits should be linked periodically. periodic=False is not supported in emu-c. operations (list): list of operations to cycle through in the digital single-qubit rotations of each layer. Defaults to [RX, RY, RX] for hea and [RX, RY] for iia. entangler (AbstractBlock): 2-qubit entangling operation. Supports CNOT, CZ, CRX, CRY, CRZ, CPHASE. Controlld rotations will have variational parameters on the rotation angles. Defaults to CNOT

For Strategy.SDAQC strategy, accepts the following: operations (list): list of operations to cycle through in the digital single-qubit rotations of each layer. Defaults to [RX, RY, RX] for hea and [RX, RY] for iia. entangler (AbstractBlock): Hamiltonian generator for the analog entangling layer. Time parameter is considered variational. Defaults to NN interaction.

For Strategy.RYDBERG strategy, accepts the following: addressable_detuning: whether to turn on the trainable semi-local addressing pattern on the detuning (n_i terms in the Hamiltonian). Defaults to True. addressable_drive: whether to turn on the trainable semi-local addressing pattern on the drive (sigma_i^x terms in the Hamiltonian). Defaults to False. tunable_phase: whether to have a tunable phase to get both sigma^x and sigma^y rotations in the drive term. If False, only a sigma^x term will be included in the drive part of the Hamiltonian generator. Defaults to False.

tag: str | None = None class-attribute instance-attribute

String to indicate the name tag of the ansatz.

Defaults to None, in which case no tag will be applied.

FeatureMapConfig(num_features=0, basis_set=BasisSet.FOURIER, reupload_scaling=ReuploadScaling.CONSTANT, feature_range=None, target_range=None, multivariate_strategy=MultivariateStrategy.PARALLEL, feature_map_strategy=Strategy.DIGITAL, param_prefix=None, num_repeats=0, operation=None, inputs=None, tag=None) dataclass

basis_set: BasisSet | dict[str, BasisSet] = BasisSet.FOURIER class-attribute instance-attribute

Basis set for feature encoding.

Takes qadence.BasisSet. Give a single BasisSet to use the same for all features. Give a dict of (str, BasisSet) where the key is the name of the variable and the value is the BasisSet to use for encoding that feature. BasisSet.FOURIER for Fourier encoding. BasisSet.CHEBYSHEV for Chebyshev encoding.

feature_map_strategy: Strategy = Strategy.DIGITAL class-attribute instance-attribute

Strategy for feature map.

Accepts DIGITAL, ANALOG or RYDBERG. Defaults to DIGITAL. If the strategy is incompatible with the operation chosen, then operation gets preference and the given strategy is ignored.

feature_range: tuple[float, float] | dict[str, tuple[float, float]] | None = None class-attribute instance-attribute

Range of data that the input data is assumed to come from.

Give a single tuple to use the same range for all features. Give a dict of (str, tuple) where the key is the name of the variable and the value is the feature range to use for that feature.

inputs: list[Basic | str] | None = None class-attribute instance-attribute

List that indicates the order of variables of the tensors that are passed.

Optional if a single feature is being encoded, required otherwise. Given input tensors xs = torch.rand(batch_size, input_size:=2) a QNN with inputs=["t", "x"] will assign t, x = xs[:,0], xs[:,1].

multivariate_strategy: MultivariateStrategy = MultivariateStrategy.PARALLEL class-attribute instance-attribute

The encoding strategy in case of multi-variate function.

Takes qadence.MultivariateStrategy. If PARALLEL, the features are encoded in one block of rotation gates with the register being split in sub-registers for each feature. If SERIES, the features are encoded sequentially using the full register for each feature, with an ansatz block between them. PARALLEL is allowed only for DIGITAL feature_map_strategy.

num_features: int = 0 class-attribute instance-attribute

Number of feature parameters to be encoded.

Defaults to 0. Thus, no feature parameters are encoded.

num_repeats: int | dict[str, int] = 0 class-attribute instance-attribute

Number of feature map layers repeated in the data reuploading step.

If all features are to be repeated the same number of times, then can give a single int. For different number of repetitions for each feature, provide a dict of (str, int) where the key is the name of the variable and the value is the number of repetitions for that feature. This amounts to the number of additional reuploads. So if num_repeats is N, the data gets uploaded N+1 times. Defaults to no repetition.

operation: Callable[[Parameter | Basic], AnalogBlock] | Type[RX] | None = None class-attribute instance-attribute

Type of operation.

Choose among the analog or digital rotations or a custom callable function returning an AnalogBlock instance. If the type of operation is incompatible with the strategy chosen, then operation gets preference and the given strategy is ignored.

param_prefix: str | None = None class-attribute instance-attribute

String prefix to create trainable parameters in Feature Map.

A string prefix to create trainable parameters multiplying the feature parameter inside the feature-encoding function. Note that currently this does not take into account the domain of the feature-encoding function. Defaults to None and thus, the feature map is not trainable. Note that this is separate from the name of the parameter. The user can provide a single prefix for all features, and it will be appended by appropriate feature name automatically.

reupload_scaling: ReuploadScaling | dict[str, ReuploadScaling] = ReuploadScaling.CONSTANT class-attribute instance-attribute

Scaling for encoding the same feature on different qubits.

Scaling used to encode the same feature on different qubits in the same layer of the feature maps. Takes qadence.ReuploadScaling. Give a single ReuploadScaling to use the same for all features. Give a dict of (str, ReuploadScaling) where the key is the name of the variable and the value is the ReuploadScaling to use for encoding that feature. ReuploadScaling.CONSTANT for constant scaling. ReuploadScaling.TOWER for linearly increasing scaling. ReuploadScaling.EXP for exponentially increasing scaling.

tag: str | None = None class-attribute instance-attribute

String to indicate the name tag of the feature map.

Defaults to None, in which case no tag will be applied.

target_range: tuple[float, float] | dict[str, tuple[float, float]] | None = None class-attribute instance-attribute

Range of data the data encoder assumes as natural range.

Give a single tuple to use the same range for all features. Give a dict of (str, tuple) where the key is the name of the variable and the value is the target range to use for that feature.

TrainConfig(max_iter=10000, print_every=0, write_every=0, checkpoint_every=0, plot_every=0, callbacks=lambda: list()(), log_model=False, root_folder=Path('./qml_logs'), create_subfolder_per_run=False, log_folder=Path('./'), checkpoint_best_only=False, val_every=0, val_epsilon=1e-05, validation_criterion=None, trainstop_criterion=None, batch_size=1, verbose=True, tracking_tool=ExperimentTrackingTool.TENSORBOARD, hyperparams=dict(), plotting_functions=tuple(), _subfolders=list()) dataclass

Default configuration for the training process.

This class provides default settings for various aspects of the training loop, such as logging, checkpointing, and validation. The default values for these fields can be customized when an instance of TrainConfig is created.

Example:

from qadence.ml_tools import TrainConfig
c = TrainConfig(root_folder="/tmp/train")

TrainConfig(max_iter=10000, print_every=0, write_every=0, checkpoint_every=0, plot_every=0, callbacks=[], log_model=False, root_folder='/tmp/train', create_subfolder_per_run=False, log_folder=PosixPath('.'), checkpoint_best_only=False, val_every=0, val_epsilon=1e-05, validation_criterion=None, trainstop_criterion=None, batch_size=1, verbose=True, tracking_tool=<ExperimentTrackingTool.TENSORBOARD: 'tensorboard'>, hyperparams={}, plotting_functions=(), _subfolders=[])

batch_size: int = 1 class-attribute instance-attribute

The batch size to use when processing a list or tuple of torch.Tensors.

This specifies how many samples are processed in each training iteration.

callbacks: list = field(default_factory=lambda: list()) class-attribute instance-attribute

List of callbacks to execute during training.

Callbacks can be used for custom behaviors, such as early stopping, custom logging, or other actions triggered at specific events.

checkpoint_best_only: bool = False class-attribute instance-attribute

If True, checkpoints are only saved if there is an improvement in the.

validation metric. This conserves storage by only keeping the best models.

validation_criterion is required when this is set to True.

checkpoint_every: int = 0 class-attribute instance-attribute

Frequency (in epochs) for saving model and optimizer checkpoints during training.

Set to 0 to disable checkpointing. This helps in resuming training or recovering models. Note that setting checkpoint_best_only = True will disable this and only best checkpoints will be saved.

create_subfolder_per_run: bool = False class-attribute instance-attribute

Whether to create a subfolder for each run, named <id>_<timestamp>_<PID>.

This ensures logs and checkpoints from different runs do not overwrite each other, which is helpful for rapid prototyping. If False, training will resume from the latest checkpoint if one exists in the specified log folder.

hyperparams: dict = field(default_factory=dict) class-attribute instance-attribute

A dictionary of hyperparameters to be tracked.

This can include learning rates, regularization parameters, or any other training-related configurations.

log_folder: Path = Path('./') class-attribute instance-attribute

The log folder for saving checkpoints and tensorboard logs.

This stores the path where all logs and checkpoints are being saved for this training session. log_folder takes precedence over root_folder and create_subfolder_per_run arguments. If the user specifies a log_folder, all checkpoints will be saved in this folder and root_folder argument will not be used.

log_model: bool = False class-attribute instance-attribute

Whether to log a serialized version of the model.

When set to True, the model's state will be logged, useful for model versioning and reproducibility.

max_iter: int = 10000 class-attribute instance-attribute

Number of training iterations (epochs) to perform.

This defines the total number of times the model will be updated.

In case of InfiniteTensorDataset, each epoch will have 1 batch. In case of TensorDataset, each epoch will have len(dataloader) batches.

plot_every: int = 0 class-attribute instance-attribute

Frequency (in epochs) for generating and saving figures during training.

Set to 0 to disable plotting.

plotting_functions: tuple[LoggablePlotFunction, ...] = field(default_factory=tuple) class-attribute instance-attribute

Functions used for in-training plotting.

These are called to generate plots that are logged or saved at specified intervals.

print_every: int = 0 class-attribute instance-attribute

Frequency (in epochs) for printing loss and metrics to the console during training.

Set to 0 to disable this output, meaning that metrics and loss will not be printed during training.

root_folder: Path = Path('./qml_logs') class-attribute instance-attribute

The root folder for saving checkpoints and tensorboard logs.

The default path is "./qml_logs"

This can be set to a specific directory where training artifacts are to be stored. Checkpoints will be saved inside a subfolder in this directory. Subfolders will be created based on create_subfolder_per_run argument.

tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD class-attribute instance-attribute

The tool used for tracking training progress and logging metrics.

Options include tools like TensorBoard, which help visualize and monitor model training.

trainstop_criterion: Callable | None = None class-attribute instance-attribute

A function to determine if the training process should stop based on a.

specific stopping metric. If None, training continues until max_iter is reached.

val_epsilon: float = 1e-05 class-attribute instance-attribute

A small safety margin used to compare the current validation loss with the.

best previous validation loss. This is used to determine improvements in metrics.

val_every: int = 0 class-attribute instance-attribute

Frequency (in epochs) for performing validation.

If set to 0, validation is not performed. Note that metrics from validation are always written, regardless of the write_every setting. Note that initial validation happens at the start of training (when val_every > 0) For initial validation - initial metrics are written. - checkpoint is saved (when checkpoint_best_only = False)

validation_criterion: Callable | None = None class-attribute instance-attribute

A function to evaluate whether a given validation metric meets a desired condition.

The validation_criterion has the following format: def validation_criterion(val_loss: float, best_val_loss: float, val_epsilon: float) -> bool: # process

If None, no custom validation criterion is applied.

verbose: bool = True class-attribute instance-attribute

Whether to print metrics and status messages during training.

If True, detailed metrics and status updates will be displayed in the console.

write_every: int = 0 class-attribute instance-attribute

Frequency (in epochs) for writing loss and metrics using the tracking tool during training.

Set to 0 to disable this logging, which prevents metrics from being logged to the tracking tool. Note that the metrics will always be written at the end of training regardless of this setting.

get_parameters(model)

Retrieve all trainable model parameters in a single vector.

PARAMETER	DESCRIPTION
model	the input PyTorch model TYPE: Module

RETURNS	DESCRIPTION
Tensor	a 1-dimensional tensor with the parameters TYPE: Tensor

Source code in qadence/ml_tools/parameters.py

def get_parameters(model: Module) -> Tensor:
    """Retrieve all trainable model parameters in a single vector.

    Args:
        model (Module): the input PyTorch model

    Returns:
        Tensor: a 1-dimensional tensor with the parameters
    """
    ps = [p.reshape(-1) for p in model.parameters() if p.requires_grad]
    return torch.concat(ps)

num_parameters(model)

Return the total number of parameters of the given model.

Source code in qadence/ml_tools/parameters.py

def num_parameters(model: Module) -> int:
    """Return the total number of parameters of the given model."""
    return len(get_parameters(model))

set_parameters(model, theta)

Set all trainable parameters of a model from a single vector.

Notice that this function assumes prior knowledge of right number of parameters in the model

PARAMETER	DESCRIPTION
model	the input PyTorch model TYPE: Module
theta	the parameters to assign TYPE: Tensor

Source code in qadence/ml_tools/parameters.py

def set_parameters(model: Module, theta: Tensor) -> None:
    """Set all trainable parameters of a model from a single vector.

    Notice that this function assumes prior knowledge of right number
    of parameters in the model

    Args:
        model (Module): the input PyTorch model
        theta (Tensor): the parameters to assign
    """

    with torch.no_grad():
        idx = 0
        for ps in model.parameters():
            if ps.requires_grad:
                n = torch.numel(ps)
                if ps.ndim == 0:
                    ps[()] = theta[idx : idx + n]
                else:
                    ps[:] = theta[idx : idx + n].reshape(ps.size())
                idx += n

optimize_step(model, optimizer, loss_fn, xs, device=None, dtype=None)

Default Torch optimize step with closure.

This is the default optimization step.

PARAMETER	DESCRIPTION
model	The input model to be optimized. TYPE: Module
optimizer	The chosen Torch optimizer. TYPE: Optimizer
loss_fn	A custom loss function that returns the loss value and a dictionary of metrics. TYPE: Callable
xs	The input data. If None, it means the given model does not require any input data. TYPE: dict | list | Tensor | None
device	A target device to run computations on. TYPE: device DEFAULT: None
dtype	Data type for xs conversion. TYPE: dtype DEFAULT: None

RETURNS	DESCRIPTION
tuple[Tensor | float, dict | None]	tuple[Tensor | float, dict | None]: A tuple containing the computed loss value and a dictionary with collected metrics.

Source code in qadence/ml_tools/optimize_step.py

def optimize_step(
    model: Module,
    optimizer: Optimizer,
    loss_fn: Callable,
    xs: dict | list | torch.Tensor | None,
    device: torch.device = None,
    dtype: torch.dtype = None,
) -> tuple[torch.Tensor | float, dict | None]:
    """Default Torch optimize step with closure.

    This is the default optimization step.

    Args:
        model (Module): The input model to be optimized.
        optimizer (Optimizer): The chosen Torch optimizer.
        loss_fn (Callable): A custom loss function
            that returns the loss value and a dictionary of metrics.
        xs (dict | list | Tensor | None): The input data. If None, it means
            the given model does not require any input data.
        device (torch.device): A target device to run computations on.
        dtype (torch.dtype): Data type for `xs` conversion.

    Returns:
        tuple[Tensor | float, dict | None]: A tuple containing the computed loss value
            and a dictionary with collected metrics.
    """

    loss, metrics = None, {}
    xs_to_device = data_to_device(xs, device=device, dtype=dtype)

    def closure() -> Any:
        # NOTE: We need the nonlocal as we can't return a metric dict and
        # because e.g. LBFGS calls this closure multiple times but for some
        # reason the returned loss is always the first one...
        nonlocal metrics, loss
        optimizer.zero_grad()
        loss, metrics = loss_fn(model, xs_to_device)
        loss.backward(retain_graph=True)
        return loss.item()

    optimizer.step(closure)
    # return the loss/metrics that are being mutated inside the closure...
    return loss, metrics

update_ng_parameters(model, optimizer, loss_fn, data, ng_params)

Update the model parameters using Nevergrad.

This function integrates Nevergrad for derivative-free optimization.

PARAMETER	DESCRIPTION
model	The PyTorch model to be optimized. TYPE: Module
optimizer	A Nevergrad optimizer instance. TYPE: Optimizer
loss_fn	A custom loss function that returns the loss value and a dictionary of metrics. TYPE: Callable[[Module, Tensor | None], tuple[float, dict]]
data	Input data for the model. If None, it means the model does not require input data. TYPE: Tensor | None
ng_params	The current set of parameters managed by Nevergrad. TYPE: Array

RETURNS	DESCRIPTION
tuple[float, dict, Array]	tuple[float, dict, ng.p.Array]: A tuple containing the computed loss value, a dictionary of metrics, and the updated Nevergrad parameters.

Source code in qadence/ml_tools/optimize_step.py

def update_ng_parameters(
    model: Module,
    optimizer: ng.optimizers.Optimizer,
    loss_fn: Callable[[Module, torch.Tensor | None], tuple[float, dict]],
    data: torch.Tensor | None,
    ng_params: ng.p.Array,
) -> tuple[float, dict, ng.p.Array]:
    """Update the model parameters using Nevergrad.

    This function integrates Nevergrad for derivative-free optimization.

    Args:
        model (Module): The PyTorch model to be optimized.
        optimizer (ng.optimizers.Optimizer): A Nevergrad optimizer instance.
        loss_fn (Callable[[Module, Tensor | None], tuple[float, dict]]): A custom loss function
            that returns the loss value and a dictionary of metrics.
        data (Tensor | None): Input data for the model. If None, it means the model does
            not require input data.
        ng_params (ng.p.Array): The current set of parameters managed by Nevergrad.

    Returns:
        tuple[float, dict, ng.p.Array]: A tuple containing the computed loss value,
            a dictionary of metrics, and the updated Nevergrad parameters.
    """
    loss, metrics = loss_fn(model, data)  # type: ignore[misc]
    optimizer.tell(ng_params, float(loss))
    ng_params = optimizer.ask()  # type: ignore[assignment]
    params = promote_to_tensor(ng_params.value, requires_grad=False)
    set_parameters(model, params)
    return loss, metrics, ng_params

DictDataLoader(dataloaders) dataclass

This class only holds a dictionary of DataLoaders and samples from them.

InfiniteTensorDataset(*tensors)

Bases: IterableDataset

Randomly sample points from the first dimension of the given tensors.

Behaves like a normal torch Dataset just that we can sample from it as many times as we want.

Examples:

import torch
from qadence.ml_tools.data import InfiniteTensorDataset

x_data, y_data = torch.rand(5,2), torch.ones(5,1)
# The dataset accepts any number of tensors with the same batch dimension
ds = InfiniteTensorDataset(x_data, y_data)

# call `next` to get one sample from each tensor:
xs = next(iter(ds))

(tensor([0.0156, 0.7122]), tensor([1.]))

Source code in qadence/ml_tools/data.py

def __init__(self, *tensors: Tensor):
    """Randomly sample points from the first dimension of the given tensors.

    Behaves like a normal torch `Dataset` just that we can sample from it as
    many times as we want.

    Examples:
    ```python exec="on" source="above" result="json"
    import torch
    from qadence.ml_tools.data import InfiniteTensorDataset

    x_data, y_data = torch.rand(5,2), torch.ones(5,1)
    # The dataset accepts any number of tensors with the same batch dimension
    ds = InfiniteTensorDataset(x_data, y_data)

    # call `next` to get one sample from each tensor:
    xs = next(iter(ds))
    print(str(xs)) # markdown-exec: hide
    ```
    """
    self.tensors = tensors
    self.indices = list(range(self.tensors[0].size(0)))

OptimizeResult(iteration, model, optimizer, loss=None, metrics=lambda: dict()(), extra=lambda: dict()()) dataclass

OptimizeResult stores many optimization intermediate values.

We store at a current iteration, the model, optimizer, loss values, metrics. An extra dict can be used for saving other information to be used for callbacks.

extra: dict = field(default_factory=lambda: dict()) class-attribute instance-attribute

Extra dict for saving anything else to be used in callbacks.

iteration: int instance-attribute

Current iteration number.

loss: Tensor | float | None = None class-attribute instance-attribute

Loss value.

metrics: dict = field(default_factory=lambda: dict()) class-attribute instance-attribute

Metrics that can be saved during training.

model: Module instance-attribute

Model at iteration.

optimizer: Optimizer | NGOptimizer instance-attribute

Optimizer at iteration.

data_to_device(xs, *args, **kwargs)

Utility method to move arbitrary data to 'device'.

Source code in qadence/ml_tools/data.py

@singledispatch
def data_to_device(xs: Any, *args: Any, **kwargs: Any) -> Any:
    """Utility method to move arbitrary data to 'device'."""
    raise ValueError(f"Unable to move {type(xs)} with input args: {args} and kwargs: {kwargs}.")

to_dataloader(*tensors, batch_size=1, infinite=False)

Convert torch tensors an (infinite) Dataloader.

PARAMETER	DESCRIPTION
*tensors	Torch tensors to use in the dataloader. TYPE: Tensor DEFAULT: ()
batch_size	batch size of sampled tensors TYPE: int DEFAULT: 1
infinite	if True, the dataloader will keep sampling indefinitely even after the whole dataset was sampled once TYPE: bool DEFAULT: False

Examples:

import torch
from qadence.ml_tools import to_dataloader

(x, y, z) = [torch.rand(10) for _ in range(3)]
loader = iter(to_dataloader(x, y, z, batch_size=5, infinite=True))
print(next(loader))
print(next(loader))
print(next(loader))

[tensor([0.6089, 0.5589, 0.6142, 0.2165, 0.6462]), tensor([0.3149, 0.6751, 0.1210, 0.8103, 0.9142]), tensor([0.2031, 0.9552, 0.8790, 0.5408, 0.5527])]
[tensor([0.5184, 0.1583, 0.2599, 0.3707, 0.7417]), tensor([0.3621, 0.0112, 0.1667, 0.2752, 0.2930]), tensor([0.1747, 0.9271, 0.7204, 0.5249, 0.1339])]
[tensor([0.6089, 0.5589, 0.6142, 0.2165, 0.6462]), tensor([0.3149, 0.6751, 0.1210, 0.8103, 0.9142]), tensor([0.2031, 0.9552, 0.8790, 0.5408, 0.5527])]

Source code in qadence/ml_tools/data.py

def to_dataloader(*tensors: Tensor, batch_size: int = 1, infinite: bool = False) -> DataLoader:
    """Convert torch tensors an (infinite) Dataloader.

    Arguments:
        *tensors: Torch tensors to use in the dataloader.
        batch_size: batch size of sampled tensors
        infinite: if `True`, the dataloader will keep sampling indefinitely even after the whole
            dataset was sampled once

    Examples:

    ```python exec="on" source="above" result="json"
    import torch
    from qadence.ml_tools import to_dataloader

    (x, y, z) = [torch.rand(10) for _ in range(3)]
    loader = iter(to_dataloader(x, y, z, batch_size=5, infinite=True))
    print(next(loader))
    print(next(loader))
    print(next(loader))
    ```
    """
    ds = InfiniteTensorDataset(*tensors) if infinite else TensorDataset(*tensors)
    return DataLoader(ds, batch_size=batch_size)

QNN(circuit, observable, backend=BackendName.PYQTORCH, diff_mode=DiffMode.AD, measurement=None, noise=None, configuration=None, inputs=None, input_diff_mode=InputDiffMode.AD)

Bases: QuantumModel

Quantum neural network model for n-dimensional inputs.

Examples:

import torch
from qadence import QuantumCircuit, QNN, Z
from qadence import hea, feature_map, hamiltonian_factory, kron

# create the circuit
n_qubits, depth = 2, 4
fm = kron(
    feature_map(1, support=(0,), param="x"),
    feature_map(1, support=(1,), param="y")
)
ansatz = hea(n_qubits=n_qubits, depth=depth)
circuit = QuantumCircuit(n_qubits, fm, ansatz)
obs_base = hamiltonian_factory(n_qubits, detuning=Z)

# the QNN will yield two outputs
obs = [2.0 * obs_base, 4.0 * obs_base]

# initialize and use the model
qnn = QNN(circuit, obs, inputs=["x", "y"])
y = qnn(torch.rand(3, 2))

tensor([[-1.5711, -3.1421],
        [-1.9156, -3.8311],
        [-1.3072, -2.6145]], grad_fn=<CatBackward0>)

Initialize the QNN.

The number of inputs is determined by the feature parameters in the input quantum circuit while the number of outputs is determined by how many observables are provided as input

PARAMETER	DESCRIPTION
circuit	The quantum circuit to use for the QNN. TYPE: QuantumCircuit
observable	The observable. TYPE: list[AbstractBlock] | AbstractBlock
backend	The chosen quantum backend. TYPE: BackendName DEFAULT: PYQTORCH
diff_mode	The differentiation engine to use. Choices 'gpsr' or 'ad'. TYPE: DiffMode DEFAULT: AD
measurement	optional measurement protocol. If None, use exact expectation value with a statevector simulator TYPE: Measurements | None DEFAULT: None
noise	A noise model to use. TYPE: NoiseHandler | None DEFAULT: None
configuration	optional configuration for the backend TYPE: BackendConfiguration | dict | None DEFAULT: None
inputs	List that indicates the order of variables of the tensors that are passed to the model. Given input tensors xs = torch.rand(batch_size, input_size:=2) a QNN with inputs=["t", "x"] will assign t, x = xs[:,0], xs[:,1]. TYPE: list[Basic | str] | None DEFAULT: None
input_diff_mode	The differentiation mode for the input tensor. TYPE: InputDiffMode | str DEFAULT: AD

Source code in qadence/ml_tools/models.py

def __init__(
    self,
    circuit: QuantumCircuit,
    observable: list[AbstractBlock] | AbstractBlock,
    backend: BackendName = BackendName.PYQTORCH,
    diff_mode: DiffMode = DiffMode.AD,
    measurement: Measurements | None = None,
    noise: NoiseHandler | None = None,
    configuration: BackendConfiguration | dict | None = None,
    inputs: list[sympy.Basic | str] | None = None,
    input_diff_mode: InputDiffMode | str = InputDiffMode.AD,
):
    """Initialize the QNN.

    The number of inputs is determined by the feature parameters in the input
    quantum circuit while the number of outputs is determined by how many
    observables are provided as input

    Args:
        circuit: The quantum circuit to use for the QNN.
        observable: The observable.
        backend: The chosen quantum backend.
        diff_mode: The differentiation engine to use. Choices 'gpsr' or 'ad'.
        measurement: optional measurement protocol. If None,
            use exact expectation value with a statevector simulator
        noise: A noise model to use.
        configuration: optional configuration for the backend
        inputs: List that indicates the order of variables of the tensors that are passed
            to the model. Given input tensors `xs = torch.rand(batch_size, input_size:=2)` a QNN
            with `inputs=["t", "x"]` will assign `t, x = xs[:,0], xs[:,1]`.
        input_diff_mode: The differentiation mode for the input tensor.
    """
    super().__init__(
        circuit,
        observable=observable,
        backend=backend,
        diff_mode=diff_mode,
        measurement=measurement,
        configuration=configuration,
        noise=noise,
    )
    if self._observable is None:
        raise ValueError("You need to provide at least one observable in the QNN constructor")
    if (inputs is not None) and (len(self.inputs) == len(inputs)):
        self.inputs = [sympy.symbols(x) if isinstance(x, str) else x for x in inputs]  # type: ignore[union-attr]
    elif (inputs is None) and len(self.inputs) <= 1:
        self.inputs = [sympy.symbols(x) if isinstance(x, str) else x for x in self.inputs]  # type: ignore[union-attr]
    else:
        raise ValueError(
            """
            Your QNN has more than one input. Please provide a list of inputs in the order of
            your tensor domain. For example, if you want to pass
            `xs = torch.rand(batch_size, input_size:=3)` to you QNN, where
            ```
            t = x[:,0]
            x = x[:,1]
            y = x[:,2]
            ```
            you have to specify
            ```
            QNN(circuit, observable, inputs=["t", "x", "y"])
            ```
            You can also pass a list of sympy symbols.
        """
        )
    self.format_to_dict = format_to_dict_fn(self.inputs)  # type: ignore[arg-type]
    self.input_diff_mode = InputDiffMode(input_diff_mode)
    if self.input_diff_mode == InputDiffMode.FD:
        from qadence.backends.utils import finitediff

        self.__derivative = finitediff
    elif self.input_diff_mode == InputDiffMode.AD:
        self.__derivative = _torch_derivative  # type: ignore[assignment]
    else:
        raise ValueError(f"Unkown forward diff mode: {self.input_diff_mode}")

forward(values=None, state=None, measurement=None, noise=None, endianness=Endianness.BIG)

Forward pass of the model.

This returns the (differentiable) expectation value of the given observable operator defined in the constructor. Differently from the base QuantumModel class, the QNN accepts also a tensor as input for the forward pass. The tensor is expected to have shape: n_batches x in_features where n_batches is the number of data points and in_features is the dimensionality of the problem

The output of the forward pass is the expectation value of the input observable(s). If a single observable is given, the output shape is n_batches while if multiple observables are given the output shape is instead n_batches x n_observables

PARAMETER	DESCRIPTION
values	the values of the feature parameters TYPE: dict[str, Tensor] | Tensor DEFAULT: None
state	Initial state. TYPE: Tensor | None DEFAULT: None
measurement	optional measurement protocol. If None, use exact expectation value with a statevector simulator TYPE: Measurements | None DEFAULT: None
noise	A noise model to use. TYPE: NoiseHandler | None DEFAULT: None
endianness	Endianness of the resulting bit strings. TYPE: Endianness DEFAULT: BIG

RETURNS	DESCRIPTION
Tensor	a tensor with the expectation value of the observables passed in the constructor of the model TYPE: Tensor

Source code in qadence/ml_tools/models.py

def forward(
    self,
    values: dict[str, Tensor] | Tensor = None,
    state: Tensor | None = None,
    measurement: Measurements | None = None,
    noise: NoiseHandler | None = None,
    endianness: Endianness = Endianness.BIG,
) -> Tensor:
    """Forward pass of the model.

    This returns the (differentiable) expectation value of the given observable
    operator defined in the constructor. Differently from the base QuantumModel
    class, the QNN accepts also a tensor as input for the forward pass. The
    tensor is expected to have shape: `n_batches x in_features` where `n_batches`
    is the number of data points and `in_features` is the dimensionality of the problem

    The output of the forward pass is the expectation value of the input
    observable(s). If a single observable is given, the output shape is
    `n_batches` while if multiple observables are given the output shape
    is instead `n_batches x n_observables`

    Args:
        values: the values of the feature parameters
        state: Initial state.
        measurement: optional measurement protocol. If None,
            use exact expectation value with a statevector simulator
        noise: A noise model to use.
        endianness: Endianness of the resulting bit strings.

    Returns:
        Tensor: a tensor with the expectation value of the observables passed
            in the constructor of the model
    """
    return self.expectation(
        values, state=state, measurement=measurement, noise=noise, endianness=endianness
    )

from_configs(register, obs_config, fm_config=FeatureMapConfig(), ansatz_config=AnsatzConfig(), backend=BackendName.PYQTORCH, diff_mode=DiffMode.AD, measurement=None, noise=None, configuration=None, input_diff_mode=InputDiffMode.AD) classmethod

Create a QNN from a set of configurations.

PARAMETER	DESCRIPTION
register	The number of qubits or a register object. TYPE: int | Register
obs_config	The configuration(s) for the observable(s). TYPE: list[ObservableConfig] | ObservableConfig
fm_config	The configuration for the feature map. Defaults to no feature encoding block. TYPE: FeatureMapConfig DEFAULT: FeatureMapConfig()
ansatz_config	The configuration for the ansatz. Defaults to a single layer of hardware efficient ansatz. TYPE: AnsatzConfig DEFAULT: AnsatzConfig()
backend	The chosen quantum backend. TYPE: BackendName DEFAULT: PYQTORCH
diff_mode	The differentiation engine to use. Choices are 'gpsr' or 'ad'. TYPE: DiffMode DEFAULT: AD
measurement	Optional measurement protocol. If None, use exact expectation value with a statevector simulator. TYPE: Measurements DEFAULT: None
noise	A noise model to use. TYPE: Noise DEFAULT: None
configuration	Optional backend configuration. TYPE: BackendConfiguration | dict DEFAULT: None
input_diff_mode	The differentiation mode for the input tensor. TYPE: InputDiffMode DEFAULT: AD

RETURNS	DESCRIPTION
QNN	A QNN object.

RAISES	DESCRIPTION
ValueError	If the observable configuration is not provided.

Example:

import torch
from qadence.ml_tools.config import AnsatzConfig, FeatureMapConfig
from qadence.ml_tools import QNN
from qadence.constructors import ObservableConfig
from qadence.operations import Z
from qadence.types import (
    AnsatzType, BackendName, BasisSet, ObservableTransform, ReuploadScaling, Strategy
)

register = 4
obs_config = ObservableConfig(
    detuning=Z,
    scale=5.0,
    shift=0.0,
    transformation_type=ObservableTransform.SCALE,
    trainable_transform=None,
)
fm_config = FeatureMapConfig(
    num_features=2,
    inputs=["x", "y"],
    basis_set=BasisSet.FOURIER,
    reupload_scaling=ReuploadScaling.CONSTANT,
    feature_range={
        "x": (-1.0, 1.0),
        "y": (0.0, 1.0),
    },
)
ansatz_config = AnsatzConfig(
    depth=2,
    ansatz_type=AnsatzType.HEA,
    ansatz_strategy=Strategy.DIGITAL,
)

qnn = QNN.from_configs(
    register, obs_config, fm_config, ansatz_config, backend=BackendName.PYQTORCH
)

x = torch.rand(2, 2)
y = qnn(x)

tensor([[1.5914],
        [2.4987]], grad_fn=<CatBackward0>)

Source code in qadence/ml_tools/models.py

@classmethod
def from_configs(
    cls,
    register: int | Register,
    obs_config: Any,
    fm_config: Any = FeatureMapConfig(),
    ansatz_config: Any = AnsatzConfig(),
    backend: BackendName = BackendName.PYQTORCH,
    diff_mode: DiffMode = DiffMode.AD,
    measurement: Measurements | None = None,
    noise: NoiseHandler | None = None,
    configuration: BackendConfiguration | dict | None = None,
    input_diff_mode: InputDiffMode | str = InputDiffMode.AD,
) -> QNN:
    """Create a QNN from a set of configurations.

    Args:
        register (int | Register): The number of qubits or a register object.
        obs_config (list[ObservableConfig] | ObservableConfig): The configuration(s)
            for the observable(s).
        fm_config (FeatureMapConfig): The configuration for the feature map.
            Defaults to no feature encoding block.
        ansatz_config (AnsatzConfig): The configuration for the ansatz.
            Defaults to a single layer of hardware efficient ansatz.
        backend (BackendName): The chosen quantum backend.
        diff_mode (DiffMode): The differentiation engine to use. Choices are
            'gpsr' or 'ad'.
        measurement (Measurements): Optional measurement protocol. If None,
            use exact expectation value with a statevector simulator.
        noise (Noise): A noise model to use.
        configuration (BackendConfiguration | dict): Optional backend configuration.
        input_diff_mode (InputDiffMode): The differentiation mode for the input tensor.

    Returns:
        A QNN object.

    Raises:
        ValueError: If the observable configuration is not provided.

    Example:
    ```python exec="on" source="material-block" result="json"
    import torch
    from qadence.ml_tools.config import AnsatzConfig, FeatureMapConfig
    from qadence.ml_tools import QNN
    from qadence.constructors import ObservableConfig
    from qadence.operations import Z
    from qadence.types import (
        AnsatzType, BackendName, BasisSet, ObservableTransform, ReuploadScaling, Strategy
    )

    register = 4
    obs_config = ObservableConfig(
        detuning=Z,
        scale=5.0,
        shift=0.0,
        transformation_type=ObservableTransform.SCALE,
        trainable_transform=None,
    )
    fm_config = FeatureMapConfig(
        num_features=2,
        inputs=["x", "y"],
        basis_set=BasisSet.FOURIER,
        reupload_scaling=ReuploadScaling.CONSTANT,
        feature_range={
            "x": (-1.0, 1.0),
            "y": (0.0, 1.0),
        },
    )
    ansatz_config = AnsatzConfig(
        depth=2,
        ansatz_type=AnsatzType.HEA,
        ansatz_strategy=Strategy.DIGITAL,
    )

    qnn = QNN.from_configs(
        register, obs_config, fm_config, ansatz_config, backend=BackendName.PYQTORCH
    )

    x = torch.rand(2, 2)
    y = qnn(x)
    print(str(y)) # markdown-exec: hide
    ```
    """
    from .constructors import build_qnn_from_configs

    return build_qnn_from_configs(
        register=register,
        observable_config=obs_config,
        fm_config=fm_config,
        ansatz_config=ansatz_config,
        backend=backend,
        diff_mode=diff_mode,
        measurement=measurement,
        noise=noise,
        configuration=configuration,
        input_diff_mode=input_diff_mode,
    )

derivative(ufa, x, derivative_indices)

Compute derivatives w.r.t.

inputs of a UFA with a single output. The derivative_indices specify which derivative(s) are computed. E.g. derivative_indices=(1,2) would compute the a second order derivative w.r.t to the indices 1 and 2 of the input tensor.

PARAMETER	DESCRIPTION
ufa	The model for which we want to compute the derivative. TYPE: Module
x	(batch_size, input_size) input tensor. TYPE: Tensor
derivative_indices	Define which derivatives to compute. TYPE: tuple

Examples: If we create a UFA with three inputs and denote the first, second, and third input with x, y, and z we can compute the following derivatives w.r.t to those inputs:

import torch
from qadence.ml_tools.models import derivative, QNN
from qadence.ml_tools.config import FeatureMapConfig, AnsatzConfig
from qadence.constructors.hamiltonians import ObservableConfig
from qadence.operations import Z

fm_config = FeatureMapConfig(num_features=3, inputs=["x", "y", "z"])
ansatz_config = AnsatzConfig()
obs_config = ObservableConfig(detuning=Z)

f = QNN.from_configs(
    register=3, obs_config=obs_config, fm_config=fm_config, ansatz_config=ansatz_config,
)
inputs = torch.rand(5,3,requires_grad=True)

# df_dx
derivative(f, inputs, (0,))

# d2f_dydz
derivative(f, inputs, (1,2))

# d3fdy2dx
derivative(f, inputs, (1,1,0))

Source code in qadence/ml_tools/models.py

def derivative(ufa: torch.nn.Module, x: Tensor, derivative_indices: tuple[int, ...]) -> Tensor:
    """Compute derivatives w.r.t.

    inputs of a UFA with a single output. The
    `derivative_indices` specify which derivative(s) are computed.  E.g.
    `derivative_indices=(1,2)` would compute the a second order derivative w.r.t
    to the indices `1` and `2` of the input tensor.

    Arguments:
        ufa: The model for which we want to compute the derivative.
        x (Tensor): (batch_size, input_size) input tensor.
        derivative_indices (tuple): Define which derivatives to compute.

    Examples:
    If we create a UFA with three inputs and denote the first, second, and third
    input with `x`, `y`, and `z` we can compute the following derivatives w.r.t
    to those inputs:
    ```py exec="on" source="material-block"
    import torch
    from qadence.ml_tools.models import derivative, QNN
    from qadence.ml_tools.config import FeatureMapConfig, AnsatzConfig
    from qadence.constructors.hamiltonians import ObservableConfig
    from qadence.operations import Z

    fm_config = FeatureMapConfig(num_features=3, inputs=["x", "y", "z"])
    ansatz_config = AnsatzConfig()
    obs_config = ObservableConfig(detuning=Z)

    f = QNN.from_configs(
        register=3, obs_config=obs_config, fm_config=fm_config, ansatz_config=ansatz_config,
    )
    inputs = torch.rand(5,3,requires_grad=True)

    # df_dx
    derivative(f, inputs, (0,))

    # d2f_dydz
    derivative(f, inputs, (1,2))

    # d3fdy2dx
    derivative(f, inputs, (1,1,0))
    ```
    """
    assert ufa.out_features == 1, "Can only call `derivative` on models with 1D output."
    return ufa._derivative(x, derivative_indices)

format_to_dict_fn(inputs=[])

Format an input tensor into the format required by the forward pass.

The tensor is assumed to have dimensions: n_batches x in_features where in_features corresponds to the number of input features of the QNN

Source code in qadence/ml_tools/models.py

def format_to_dict_fn(
    inputs: list[sympy.Symbol | str] = [],
) -> Callable[[Tensor | ParamDictType], ParamDictType]:
    """Format an input tensor into the format required by the forward pass.

    The tensor is assumed to have dimensions: n_batches x in_features where in_features
    corresponds to the number of input features of the QNN
    """
    in_features = len(inputs)

    def tensor_to_dict(values: Tensor | ParamDictType) -> ParamDictType:
        if isinstance(values, Tensor):
            values = values.reshape(-1, 1) if len(values.size()) == 1 else values
            if not values.shape[1] == in_features:
                raise ValueError(
                    f"Model expects in_features={in_features} but got {values.shape[1]}."
                )
            values = {fparam.name: values[:, inputs.index(fparam)] for fparam in inputs}  # type: ignore[union-attr]
        return values

    return tensor_to_dict

Callback(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)

Base class for defining various training callbacks.

ATTRIBUTE	DESCRIPTION
on	The event on which to trigger the callback. Must be a valid on value from: ["train_start", "train_end", "train_epoch_start", "train_epoch_end", "train_batch_start", "train_batch_end","val_epoch_start", "val_epoch_end", "val_batch_start", "val_batch_end", "test_batch_start", "test_batch_end"] TYPE: str
called_every	Frequency of callback calls in terms of iterations. TYPE: int
callback	The function to call if the condition is met. TYPE: CallbackFunction | None
callback_condition	Condition to check before calling. TYPE: CallbackConditionFunction | None
modify_optimize_result	Function to modify OptimizeResult. TYPE: CallbackFunction | dict[str, Any] | None

A callback can be defined in two ways:

1. By providing a callback function directly in the base class: This is useful for simple callbacks that don't require subclassing.

Example:

from qadence.ml_tools.callbacks import Callback

def custom_callback_function(trainer, config, writer):
    print("Custom callback executed.")

custom_callback = Callback(
    on="train_end",
    called_every=5,
    callback=custom_callback_function
)

1. By inheriting and implementing the run_callback method: This is suitable for more complex callbacks that require customization.

Example:

from qadence.ml_tools.callbacks import Callback
class CustomCallback(Callback):
    def run_callback(self, trainer, config, writer):
        print("Custom behavior in the inherited run_callback method.")

custom_callback = CustomCallback(on="train_end", called_every=10)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = callback_condition or (lambda _: True)

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = (
            lambda opt_res: opt_res.extra.update(modify_optimize_result) or opt_res
        )
    else:
        self.modify_optimize_result = modify_optimize_result or (lambda opt_res: opt_res)

on: TrainingStage | str property writable

Returns the TrainingStage.

RETURNS	DESCRIPTION
TrainingStage	TrainingStage for the callback TYPE: TrainingStage | str

__call__(when, trainer, config, writer)

Executes the callback if conditions are met.

PARAMETER	DESCRIPTION
when	The event when the callback is triggered. TYPE: str
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

RETURNS	DESCRIPTION
Any	Result of the callback function if executed. TYPE: Any

Source code in qadence/ml_tools/callbacks/callback.py

def __call__(
    self, when: TrainingStage, trainer: Any, config: TrainConfig, writer: BaseWriter
) -> Any:
    """Executes the callback if conditions are met.

    Args:
        when (str): The event when the callback is triggered.
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.

    Returns:
        Any: Result of the callback function if executed.
    """
    opt_result = trainer.opt_result
    if self.on == when:
        if opt_result:
            opt_result = self.modify_optimize_result(opt_result)
        if self._should_call(when, opt_result):
            return self.run_callback(trainer, config, writer)

run_callback(trainer, config, writer)

Executes the defined callback.

PARAMETER	DESCRIPTION
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

RETURNS	DESCRIPTION
Any	Result of the callback execution. TYPE: Any

RAISES	DESCRIPTION
NotImplementedError	If not implemented in subclasses.

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Executes the defined callback.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.

    Returns:
        Any: Result of the callback execution.

    Raises:
        NotImplementedError: If not implemented in subclasses.
    """
    if self.callback is not None:
        return self.callback(trainer, config, writer)
    raise NotImplementedError("Subclasses should override the run_callback method.")

LoadCheckpoint(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)

Bases: Callback

Callback to load a model checkpoint.

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = callback_condition or (lambda _: True)

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = (
            lambda opt_res: opt_res.extra.update(modify_optimize_result) or opt_res
        )
    else:
        self.modify_optimize_result = modify_optimize_result or (lambda opt_res: opt_res)

run_callback(trainer, config, writer)

Loads a model checkpoint.

PARAMETER	DESCRIPTION
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

RETURNS	DESCRIPTION
Any	The result of loading the checkpoint. TYPE: Any

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Loads a model checkpoint.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.

    Returns:
        Any: The result of loading the checkpoint.
    """
    folder = config.log_folder
    model = trainer.model
    optimizer = trainer.optimizer
    device = trainer.log_device
    return load_checkpoint(folder, model, optimizer, device=device)

LogHyperparameters(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)

Bases: Callback

Callback to log hyperparameters using the writer.

The LogHyperparameters callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use LogHyperparameters, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import LogHyperparameters

# Create an instance of the LogHyperparameters callback
log_hyper_callback = LogHyperparameters(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[log_hyper_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = callback_condition or (lambda _: True)

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = (
            lambda opt_res: opt_res.extra.update(modify_optimize_result) or opt_res
        )
    else:
        self.modify_optimize_result = modify_optimize_result or (lambda opt_res: opt_res)

run_callback(trainer, config, writer)

Logs hyperparameters using the writer.

PARAMETER	DESCRIPTION
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Logs hyperparameters using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    hyperparams = config.hyperparams
    writer.log_hyperparams(hyperparams)

LogModelTracker(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)

Bases: Callback

Callback to log the model using the writer.

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = callback_condition or (lambda _: True)

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = (
            lambda opt_res: opt_res.extra.update(modify_optimize_result) or opt_res
        )
    else:
        self.modify_optimize_result = modify_optimize_result or (lambda opt_res: opt_res)

run_callback(trainer, config, writer)

Logs the model using the writer.

PARAMETER	DESCRIPTION
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Logs the model using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    model = trainer.model
    writer.log_model(
        model, trainer.train_dataloader, trainer.val_dataloader, trainer.test_dataloader
    )

PlotMetrics(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)

Bases: Callback

Callback to plot metrics using the writer.

The PlotMetrics callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use PlotMetrics, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import PlotMetrics

# Create an instance of the PlotMetrics callback
plot_metrics_callback = PlotMetrics(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[plot_metrics_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = callback_condition or (lambda _: True)

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = (
            lambda opt_res: opt_res.extra.update(modify_optimize_result) or opt_res
        )
    else:
        self.modify_optimize_result = modify_optimize_result or (lambda opt_res: opt_res)

run_callback(trainer, config, writer)

Plots metrics using the writer.

PARAMETER	DESCRIPTION
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Plots metrics using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    opt_result = trainer.opt_result
    plotting_functions = config.plotting_functions
    writer.plot(trainer.model, opt_result.iteration, plotting_functions)

PrintMetrics(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)

Bases: Callback

Callback to print metrics using the writer.

The PrintMetrics callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use PrintMetrics, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import PrintMetrics

# Create an instance of the PrintMetrics callback
print_metrics_callback = PrintMetrics(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[print_metrics_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = callback_condition or (lambda _: True)

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = (
            lambda opt_res: opt_res.extra.update(modify_optimize_result) or opt_res
        )
    else:
        self.modify_optimize_result = modify_optimize_result or (lambda opt_res: opt_res)

run_callback(trainer, config, writer)

Prints metrics using the writer.

PARAMETER	DESCRIPTION
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Prints metrics using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    opt_result = trainer.opt_result
    writer.print_metrics(opt_result)

SaveBestCheckpoint(on, called_every)

Bases: SaveCheckpoint

Callback to save the best model checkpoint based on a validation criterion.

Initializes the SaveBestCheckpoint callback.

PARAMETER	DESCRIPTION
on	The event to trigger the callback. TYPE: str
called_every	Frequency of callback calls in terms of iterations. TYPE: int

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(self, on: str, called_every: int):
    """Initializes the SaveBestCheckpoint callback.

    Args:
        on (str): The event to trigger the callback.
        called_every (int): Frequency of callback calls in terms of iterations.
    """
    super().__init__(on=on, called_every=called_every)
    self.best_loss = float("inf")

run_callback(trainer, config, writer)

Saves the checkpoint if the current loss is better than the best loss.

PARAMETER	DESCRIPTION
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Saves the checkpoint if the current loss is better than the best loss.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    opt_result = trainer.opt_result
    if config.validation_criterion and config.validation_criterion(
        opt_result.loss, self.best_loss, config.val_epsilon
    ):
        self.best_loss = opt_result.loss

        folder = config.log_folder
        model = trainer.model
        optimizer = trainer.optimizer
        opt_result = trainer.opt_result
        write_checkpoint(folder, model, optimizer, "best")

SaveCheckpoint(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)

Bases: Callback

Callback to save a model checkpoint.

The SaveCheckpoint callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use SaveCheckpoint, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import SaveCheckpoint

# Create an instance of the SaveCheckpoint callback
save_checkpoint_callback = SaveCheckpoint(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[save_checkpoint_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = callback_condition or (lambda _: True)

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = (
            lambda opt_res: opt_res.extra.update(modify_optimize_result) or opt_res
        )
    else:
        self.modify_optimize_result = modify_optimize_result or (lambda opt_res: opt_res)

run_callback(trainer, config, writer)

Saves a model checkpoint.

PARAMETER	DESCRIPTION
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Saves a model checkpoint.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    folder = config.log_folder
    model = trainer.model
    optimizer = trainer.optimizer
    opt_result = trainer.opt_result
    write_checkpoint(folder, model, optimizer, opt_result.iteration)

WriteMetrics(on='idle', called_every=1, callback=None, callback_condition=None, modify_optimize_result=None)

Bases: Callback

Callback to write metrics using the writer.

The WriteMetrics callback can be added to the TrainConfig callbacks as a custom user defined callback.

Example Usage in TrainConfig: To use WriteMetrics, include it in the callbacks list when setting up your TrainConfig:

from qadence.ml_tools import TrainConfig
from qadence.ml_tools.callbacks import WriteMetrics

# Create an instance of the WriteMetrics callback
write_metrics_callback = WriteMetrics(on = "val_batch_end", called_every = 100)

config = TrainConfig(
    max_iter=10000,
    # Print metrics every 1000 training epochs
    print_every=1000,
    # Add the custom callback that runs every 100 val_batch_end
    callbacks=[write_metrics_callback]
)

Source code in qadence/ml_tools/callbacks/callback.py

def __init__(
    self,
    on: str | TrainingStage = "idle",
    called_every: int = 1,
    callback: CallbackFunction | None = None,
    callback_condition: CallbackConditionFunction | None = None,
    modify_optimize_result: CallbackFunction | dict[str, Any] | None = None,
):
    if not isinstance(called_every, int):
        raise ValueError("called_every must be a positive integer or 0")

    self.callback: CallbackFunction | None = callback
    self.on: str | TrainingStage = on
    self.called_every: int = called_every
    self.callback_condition = callback_condition or (lambda _: True)

    if isinstance(modify_optimize_result, dict):
        self.modify_optimize_result = (
            lambda opt_res: opt_res.extra.update(modify_optimize_result) or opt_res
        )
    else:
        self.modify_optimize_result = modify_optimize_result or (lambda opt_res: opt_res)

run_callback(trainer, config, writer)

Writes metrics using the writer.

PARAMETER	DESCRIPTION
trainer	The training object. TYPE: Any
config	The configuration object. TYPE: TrainConfig
writer	The writer object for logging. TYPE: BaseWriter

Source code in qadence/ml_tools/callbacks/callback.py

def run_callback(self, trainer: Any, config: TrainConfig, writer: BaseWriter) -> Any:
    """Writes metrics using the writer.

    Args:
        trainer (Any): The training object.
        config (TrainConfig): The configuration object.
        writer (BaseWriter ): The writer object for logging.
    """
    opt_result = trainer.opt_result
    writer.write(opt_result)

BaseTrainer(model, optimizer, config, loss_fn='mse', optimize_step=optimize_step, train_dataloader=None, val_dataloader=None, test_dataloader=None, max_batches=None)

Base class for training machine learning models using a given optimizer.

The base class implements contextmanager for gradient based/free optimization, properties, property setters, input validations, callback decorator generator, and empty hooks for different training steps.

This class provides

• Context managers for enabling/disabling gradient-based optimization
• Properties for managing models, optimizers, and dataloaders
• Input validations and a callback decorator generator
• Config and callback managers using the provided TrainConfig

ATTRIBUTE	DESCRIPTION
use_grad	Indicates if gradients are used for optimization. Default is True. TYPE: bool
model	The neural network model. TYPE: Module
optimizer	The optimizer for training. TYPE: Optimizer | Optimizer | None
config	The configuration settings for training. TYPE: TrainConfig
train_dataloader	DataLoader for training data. TYPE: Dataloader | DictDataLoader | None
val_dataloader	DataLoader for validation data. TYPE: Dataloader | DictDataLoader | None
test_dataloader	DataLoader for testing data. TYPE: Dataloader | DictDataLoader | None
optimize_step	Function for performing an optimization step. TYPE: Callable
loss_fn	loss function to use. Default loss function used is 'mse' TYPE: Callable | str ]
num_training_batches	Number of training batches. In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: int
num_validation_batches	Number of validation batches. In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: int
num_test_batches	Number of test batches. In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: int
state	Current state in the training process TYPE: str

Initializes the BaseTrainer.

PARAMETER	DESCRIPTION
model	The model to train. TYPE: Module
optimizer	The optimizer for training. TYPE: Optimizer | Optimizer | None
config	The TrainConfig settings for training. TYPE: TrainConfig
loss_fn	The loss function to use. str input to be specified to use a default loss function. currently supported loss functions: 'mse', 'cross_entropy'. If not specified, default mse loss will be used. TYPE: str | Callable DEFAULT: 'mse'
train_dataloader	DataLoader for training data. If the model does not need data to evaluate loss, no dataset should be provided. TYPE: Dataloader | DictDataLoader | None DEFAULT: None
val_dataloader	DataLoader for validation data. TYPE: Dataloader | DictDataLoader | None DEFAULT: None
test_dataloader	DataLoader for testing data. TYPE: Dataloader | DictDataLoader | None DEFAULT: None
max_batches	Maximum number of batches to process per epoch. This is only valid in case of finite TensorDataset dataloaders. if max_batches is not None, the maximum number of batches used will be min(max_batches, len(dataloader.dataset)) In case of InfiniteTensorDataset only 1 batch per epoch is used. TYPE: int | None DEFAULT: None

Source code in qadence/ml_tools/train_utils/base_trainer.py

def __init__(
    self,
    model: nn.Module,
    optimizer: optim.Optimizer | NGOptimizer | None,
    config: TrainConfig,
    loss_fn: str | Callable = "mse",
    optimize_step: Callable = optimize_step,
    train_dataloader: DataLoader | DictDataLoader | None = None,
    val_dataloader: DataLoader | DictDataLoader | None = None,
    test_dataloader: DataLoader | DictDataLoader | None = None,
    max_batches: int | None = None,
):
    """
    Initializes the BaseTrainer.

    Args:
        model (nn.Module): The model to train.
        optimizer (optim.Optimizer | NGOptimizer | None): The optimizer
            for training.
        config (TrainConfig): The TrainConfig settings for training.
        loss_fn (str | Callable): The loss function to use.
            str input to be specified to use a default loss function.
            currently supported loss functions: 'mse', 'cross_entropy'.
            If not specified, default mse loss will be used.
        train_dataloader (Dataloader | DictDataLoader | None): DataLoader for training data.
            If the model does not need data to evaluate loss, no dataset
            should be provided.
        val_dataloader (Dataloader | DictDataLoader | None): DataLoader for validation data.
        test_dataloader (Dataloader | DictDataLoader | None): DataLoader for testing data.
        max_batches (int | None): Maximum number of batches to process per epoch.
            This is only valid in case of finite TensorDataset dataloaders.
            if max_batches is not None, the maximum number of batches used will
            be min(max_batches, len(dataloader.dataset))
            In case of InfiniteTensorDataset only 1 batch per epoch is used.
    """
    self._model: nn.Module
    self._optimizer: optim.Optimizer | NGOptimizer | None
    self._config: TrainConfig
    self._train_dataloader: DataLoader | DictDataLoader | None = None
    self._val_dataloader: DataLoader | DictDataLoader | None = None
    self._test_dataloader: DataLoader | DictDataLoader | None = None

    self.config = config
    self.model = model
    self.optimizer = optimizer
    self.max_batches = max_batches

    self.num_training_batches: int
    self.num_validation_batches: int
    self.num_test_batches: int

    self.train_dataloader = train_dataloader
    self.val_dataloader = val_dataloader
    self.test_dataloader = test_dataloader

    self.loss_fn: Callable = get_loss_fn(loss_fn)
    self.optimize_step: Callable = optimize_step
    self.ng_params: ng.p.Array
    self.training_stage: TrainingStage = TrainingStage("idle")

config: TrainConfig property writable

Returns the training configuration.

RETURNS	DESCRIPTION
TrainConfig	The configuration object. TYPE: TrainConfig

model: nn.Module property writable

Returns the model if set, otherwise raises an error.

RETURNS	DESCRIPTION
Module	nn.Module: The model.

optimizer: optim.Optimizer | NGOptimizer | None property writable

Returns the optimizer if set, otherwise raises an error.

RETURNS	DESCRIPTION
Optimizer | Optimizer | None	optim.Optimizer | NGOptimizer | None: The optimizer.

test_dataloader: DataLoader property writable

Returns the test DataLoader, validating its type.

RETURNS	DESCRIPTION
DataLoader	The DataLoader for testing data. TYPE: DataLoader

train_dataloader: DataLoader property writable

Returns the training DataLoader, validating its type.

RETURNS	DESCRIPTION
DataLoader	The DataLoader for training data. TYPE: DataLoader

use_grad: bool property writable

Returns the optimization framework for the trainer.

use_grad = True : Gradient based optimization use_grad = False : Gradient free optimization

RETURNS	DESCRIPTION
bool	Bool value for using gradient. TYPE: bool

val_dataloader: DataLoader property writable

Returns the validation DataLoader, validating its type.

RETURNS	DESCRIPTION
DataLoader	The DataLoader for validation data. TYPE: DataLoader

callback(phase) staticmethod

Decorator for executing callbacks before and after a phase.

Phase are different hooks during the training. list of valid phases is defined in Callbacks. We also update the current state of the training process in the callback decorator.

PARAMETER	DESCRIPTION
phase	The phase for which the callback is executed (e.g., "train", "train_epoch", "train_batch"). TYPE: str

RETURNS	DESCRIPTION
Callable	The decorated function. TYPE: Callable

Source code in qadence/ml_tools/train_utils/base_trainer.py

@staticmethod
def callback(phase: str) -> Callable:
    """
    Decorator for executing callbacks before and after a phase.

    Phase are different hooks during the training. list of valid
    phases is defined in Callbacks.
    We also update the current state of the training process in
    the callback decorator.

    Args:
        phase (str): The phase for which the callback is executed (e.g., "train",
            "train_epoch", "train_batch").

    Returns:
        Callable: The decorated function.
    """

    def decorator(method: Callable) -> Callable:
        def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
            start_event = f"{phase}_start"
            end_event = f"{phase}_end"

            self.training_stage = TrainingStage(start_event)
            self.callback_manager.run_callbacks(trainer=self)
            result = method(self, *args, **kwargs)

            self.training_stage = TrainingStage(end_event)
            # build_optimize_result method is defined in the trainer.
            self.build_optimize_result(result)
            self.callback_manager.run_callbacks(trainer=self)

            return result

        return wrapper

    return decorator

disable_grad_opt(optimizer=None)

Context manager to temporarily disable gradient-based optimization.

PARAMETER	DESCRIPTION
optimizer	The Nevergrad optimizer to use. If no optimizer is provided, default optimizer for trainer object will be used. TYPE: Optimizer DEFAULT: None

Source code in qadence/ml_tools/train_utils/base_trainer.py

@contextmanager
def disable_grad_opt(self, optimizer: NGOptimizer | None = None) -> Iterator[None]:
    """
    Context manager to temporarily disable gradient-based optimization.

    Args:
        optimizer (NGOptimizer): The Nevergrad optimizer to use.
            If no optimizer is provided, default optimizer for trainer
            object will be used.
    """
    original_mode = self.use_grad
    original_optimizer = self._optimizer
    try:
        self.use_grad = False
        self.callback_manager.use_grad = False
        self.optimizer = optimizer if optimizer else self.optimizer
        yield
    finally:
        self.use_grad = original_mode
        self.callback_manager.use_grad = original_mode
        self.optimizer = original_optimizer

enable_grad_opt(optimizer=None)

Context manager to temporarily enable gradient-based optimization.

PARAMETER	DESCRIPTION
optimizer	The PyTorch optimizer to use. If no optimizer is provided, default optimizer for trainer object will be used. TYPE: Optimizer DEFAULT: None

Source code in qadence/ml_tools/train_utils/base_trainer.py

@contextmanager
def enable_grad_opt(self, optimizer: optim.Optimizer | None = None) -> Iterator[None]:
    """
    Context manager to temporarily enable gradient-based optimization.

    Args:
        optimizer (optim.Optimizer): The PyTorch optimizer to use.
            If no optimizer is provided, default optimizer for trainer
            object will be used.
    """
    original_mode = self.use_grad
    original_optimizer = self._optimizer
    try:
        self.use_grad = True
        self.callback_manager.use_grad = True
        self.optimizer = optimizer if optimizer else self.optimizer
        yield
    finally:
        self.use_grad = original_mode
        self.callback_manager.use_grad = original_mode
        self.optimizer = original_optimizer

on_test_batch_end(test_batch_loss_metrics)

Called at the end of each testing batch.

PARAMETER	DESCRIPTION
test_batch_loss_metrics	Metrics for the testing batch loss. tuple of (loss, metrics) TYPE: tuple[Tensor, Any]

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_test_batch_end(self, test_batch_loss_metrics: tuple[torch.Tensor, Any]) -> None:
    """
    Called at the end of each testing batch.

    Args:
        test_batch_loss_metrics: Metrics for the testing batch loss.
            tuple of (loss, metrics)
    """
    pass

on_test_batch_start(batch)

Called at the start of each testing batch.

PARAMETER	DESCRIPTION
batch	A batch of data from the DataLoader. Typically a tuple containing input tensors and corresponding target tensors. TYPE: tuple[Tensor, ...] | None

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_test_batch_start(self, batch: tuple[torch.Tensor, ...] | None) -> None:
    """
    Called at the start of each testing batch.

    Args:
        batch: A batch of data from the DataLoader. Typically a tuple containing
            input tensors and corresponding target tensors.
    """
    pass

on_train_batch_end(train_batch_loss_metrics)

Called at the end of each training batch.

PARAMETER	DESCRIPTION
train_batch_loss_metrics	Metrics for the training batch loss. tuple of (loss, metrics) TYPE: tuple[Tensor, Any]

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_batch_end(self, train_batch_loss_metrics: tuple[torch.Tensor, Any]) -> None:
    """
    Called at the end of each training batch.

    Args:
        train_batch_loss_metrics: Metrics for the training batch loss.
            tuple of (loss, metrics)
    """
    pass

on_train_batch_start(batch)

Called at the start of each training batch.

PARAMETER	DESCRIPTION
batch	A batch of data from the DataLoader. Typically a tuple containing input tensors and corresponding target tensors. TYPE: tuple[Tensor, ...] | None

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_batch_start(self, batch: tuple[torch.Tensor, ...] | None) -> None:
    """
    Called at the start of each training batch.

    Args:
        batch: A batch of data from the DataLoader. Typically a tuple containing
            input tensors and corresponding target tensors.
    """
    pass

on_train_end(train_losses, val_losses=None)

Called at the end of training.

PARAMETER	DESCRIPTION
train_losses	Metrics for the training losses. list -> list -> tuples Epochs -> Training Batches -> (loss, metrics) TYPE: list[list[tuple[Tensor, Any]]]
val_losses	Metrics for the validation losses. list -> list -> tuples Epochs -> Validation Batches -> (loss, metrics) TYPE: list[list[tuple[Tensor, Any]]] | None DEFAULT: None

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_end(
    self,
    train_losses: list[list[tuple[torch.Tensor, Any]]],
    val_losses: list[list[tuple[torch.Tensor, Any]]] | None = None,
) -> None:
    """
    Called at the end of training.

    Args:
        train_losses (list[list[tuple[torch.Tensor, Any]]]):
            Metrics for the training losses.
            list    -> list                  -> tuples
            Epochs  -> Training Batches      -> (loss, metrics)
        val_losses (list[list[tuple[torch.Tensor, Any]]] | None):
            Metrics for the validation losses.
            list    -> list                  -> tuples
            Epochs  -> Validation Batches    -> (loss, metrics)
    """
    pass

on_train_epoch_end(train_epoch_loss_metrics)

Called at the end of each training epoch.

PARAMETER	DESCRIPTION
train_epoch_loss_metrics	Metrics for the training epoch losses. list -> tuples Training Batches -> (loss, metrics) TYPE: list[tuple[Tensor, Any]]

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_epoch_end(self, train_epoch_loss_metrics: list[tuple[torch.Tensor, Any]]) -> None:
    """
    Called at the end of each training epoch.

    Args:
        train_epoch_loss_metrics: Metrics for the training epoch losses.
            list                  -> tuples
            Training Batches      -> (loss, metrics)
    """
    pass

on_train_epoch_start()

Called at the start of each training epoch.

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_epoch_start(self) -> None:
    """Called at the start of each training epoch."""
    pass

on_train_start()

Called at the start of training.

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_train_start(self) -> None:
    """Called at the start of training."""
    pass

on_val_batch_end(val_batch_loss_metrics)

Called at the end of each validation batch.

PARAMETER	DESCRIPTION
val_batch_loss_metrics	Metrics for the validation batch loss. tuple of (loss, metrics) TYPE: tuple[Tensor, Any]

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_val_batch_end(self, val_batch_loss_metrics: tuple[torch.Tensor, Any]) -> None:
    """
    Called at the end of each validation batch.

    Args:
        val_batch_loss_metrics: Metrics for the validation batch loss.
            tuple of (loss, metrics)
    """
    pass

on_val_batch_start(batch)

Called at the start of each validation batch.

PARAMETER	DESCRIPTION
batch	A batch of data from the DataLoader. Typically a tuple containing input tensors and corresponding target tensors. TYPE: tuple[Tensor, ...] | None

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_val_batch_start(self, batch: tuple[torch.Tensor, ...] | None) -> None:
    """
    Called at the start of each validation batch.

    Args:
        batch: A batch of data from the DataLoader. Typically a tuple containing
            input tensors and corresponding target tensors.
    """
    pass

on_val_epoch_end(val_epoch_loss_metrics)

Called at the end of each validation epoch.

PARAMETER	DESCRIPTION
val_epoch_loss_metrics	Metrics for the validation epoch loss. list -> tuples Validation Batches -> (loss, metrics) TYPE: list[tuple[Tensor, Any]]

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_val_epoch_end(self, val_epoch_loss_metrics: list[tuple[torch.Tensor, Any]]) -> None:
    """
    Called at the end of each validation epoch.

    Args:
        val_epoch_loss_metrics: Metrics for the validation epoch loss.
            list                    -> tuples
            Validation Batches      -> (loss, metrics)
    """
    pass

on_val_epoch_start()

Called at the start of each validation epoch.

Source code in qadence/ml_tools/train_utils/base_trainer.py

def on_val_epoch_start(self) -> None:
    """Called at the start of each validation epoch."""
    pass

set_use_grad(value) classmethod

Sets the global use_grad flag.

PARAMETER	DESCRIPTION
value	Whether to use gradient-based optimization. TYPE: bool

Source code in qadence/ml_tools/train_utils/base_trainer.py

@classmethod
def set_use_grad(cls, value: bool) -> None:
    """
    Sets the global use_grad flag.

    Args:
        value (bool): Whether to use gradient-based optimization.
    """
    if not isinstance(value, bool):
        raise TypeError("use_grad must be a boolean value.")
    cls._use_grad = value

BaseWriter

Bases: ABC

Abstract base class for experiment tracking writers.

METHOD	DESCRIPTION
open	Opens the writer and sets up the logging environment.
close	Closes the writer and finalizes any ongoing logging processes.
print_metrics	Prints metrics and loss in a formatted manner.
write	Writes the optimization results to the tracking tool.
log_hyperparams	Logs the hyperparameters to the tracking tool.
plot	Logs model plots using provided plotting functions.
log_model	Logs the model and any relevant information.

close() abstractmethod

Closes the writer and finalizes logging.

Source code in qadence/ml_tools/callbacks/writer_registry.py

@abstractmethod
def close(self) -> None:
    """Closes the writer and finalizes logging."""
    raise NotImplementedError("Writers must implement a close method.")

log_hyperparams(hyperparams) abstractmethod

Logs hyperparameters.

PARAMETER	DESCRIPTION
hyperparams	A dictionary of hyperparameters to log. TYPE: dict

Source code in qadence/ml_tools/callbacks/writer_registry.py

@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.")

log_model(model, train_dataloader=None, val_dataloader=None, test_dataloader=None) abstractmethod

Logs the model and associated data.

PARAMETER	DESCRIPTION
model	The model to log. TYPE: Module
train_dataloader	DataLoader for training data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
val_dataloader	DataLoader for validation data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
test_dataloader	DataLoader for testing data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None

Source code in qadence/ml_tools/callbacks/writer_registry.py

@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.")

open(config, iteration=None) abstractmethod

Opens the writer and prepares it for logging.

PARAMETER	DESCRIPTION
config	Configuration object containing settings for logging. TYPE: TrainConfig
iteration	The iteration step to start logging from. Defaults to None. TYPE: int DEFAULT: None

Source code in qadence/ml_tools/callbacks/writer_registry.py

@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.")

plot(model, iteration, plotting_functions) abstractmethod

Logs plots of the model using provided plotting functions.

PARAMETER	DESCRIPTION
model	The model to plot. TYPE: Module
iteration	The current iteration number. TYPE: int
plotting_functions	Functions used to generate plots. TYPE: tuple[PlottingFunction, ...]

Source code in qadence/ml_tools/callbacks/writer_registry.py

@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.")

print_metrics(result)

Prints the metrics and loss in a readable format.

PARAMETER	DESCRIPTION
result	The optimization results to display. TYPE: OptimizeResult

Source code in qadence/ml_tools/callbacks/writer_registry.py

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)

write(result) abstractmethod

Logs the results of the current iteration.

PARAMETER	DESCRIPTION
result	The optimization results to log. TYPE: OptimizeResult

Source code in qadence/ml_tools/callbacks/writer_registry.py

@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.")

MLFlowWriter()

Bases: BaseWriter

Writer for logging to MLflow.

ATTRIBUTE	DESCRIPTION
run	The active MLflow run. TYPE: Run
mlflow	The MLflow module. TYPE: ModuleType

Source code in qadence/ml_tools/callbacks/writer_registry.py

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

close()

Closes the MLflow run.

Source code in qadence/ml_tools/callbacks/writer_registry.py

def close(self) -> None:
    """Closes the MLflow run."""
    if self.run:
        self.mlflow.end_run()

get_signature_from_dataloader(model, dataloader)

Infers the signature of the model based on the input data from the dataloader.

PARAMETER	DESCRIPTION
model	The model to use for inference. TYPE: Module
dataloader	DataLoader for model inputs. TYPE: DataLoader | DictDataLoader | None

RETURNS	DESCRIPTION
Any	Optional[Any]: The inferred signature, if available.

Source code in qadence/ml_tools/callbacks/writer_registry.py

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

log_hyperparams(hyperparams)

Logs hyperparameters to MLflow.

PARAMETER	DESCRIPTION
hyperparams	A dictionary of hyperparameters to log. TYPE: dict

Source code in qadence/ml_tools/callbacks/writer_registry.py

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"
        )

log_model(model, train_dataloader=None, val_dataloader=None, test_dataloader=None)

Logs the model and its signature to MLflow using the provided data loaders.

PARAMETER	DESCRIPTION
model	The model to log. TYPE: Module
train_dataloader	DataLoader for training data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
val_dataloader	DataLoader for validation data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
test_dataloader	DataLoader for testing data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None

Source code in qadence/ml_tools/callbacks/writer_registry.py

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)

open(config, iteration=None)

Opens the MLflow writer and initializes an MLflow run.

PARAMETER	DESCRIPTION
config	Configuration object containing settings for logging. TYPE: TrainConfig
iteration	The iteration step to start logging from. Defaults to None. TYPE: int DEFAULT: None

RETURNS	DESCRIPTION
mlflow	The MLflow module instance. TYPE: ModuleType | None

Source code in qadence/ml_tools/callbacks/writer_registry.py

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

plot(model, iteration, plotting_functions)

Logs plots of the model using provided plotting functions.

PARAMETER	DESCRIPTION
model	The model to plot. TYPE: Module
iteration	The current iteration number. TYPE: int
plotting_functions	Functions used to generate plots. TYPE: tuple[PlottingFunction, ...]

Source code in qadence/ml_tools/callbacks/writer_registry.py

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"
        )

write(result)

Logs the results of the current iteration to MLflow.

PARAMETER	DESCRIPTION
result	The optimization results to log. TYPE: OptimizeResult

Source code in qadence/ml_tools/callbacks/writer_registry.py

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"
        )

TensorBoardWriter()

Bases: BaseWriter

Writer for logging to TensorBoard.

ATTRIBUTE	DESCRIPTION
writer	The TensorBoard SummaryWriter instance. TYPE: SummaryWriter

Source code in qadence/ml_tools/callbacks/writer_registry.py

def __init__(self) -> None:
    self.writer = None

close()

Closes the TensorBoard writer.

Source code in qadence/ml_tools/callbacks/writer_registry.py

def close(self) -> None:
    """Closes the TensorBoard writer."""
    if self.writer:
        self.writer.close()

log_hyperparams(hyperparams)

Logs hyperparameters to TensorBoard.

PARAMETER	DESCRIPTION
hyperparams	A dictionary of hyperparameters to log. TYPE: dict

Source code in qadence/ml_tools/callbacks/writer_registry.py

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"
        )

log_model(model, train_dataloader=None, val_dataloader=None, test_dataloader=None)

Logs the model.

Currently not supported by TensorBoard.

PARAMETER	DESCRIPTION
model	The model to log. TYPE: Module
train_dataloader	DataLoader for training data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
val_dataloader	DataLoader for validation data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None
test_dataloader	DataLoader for testing data. TYPE: DataLoader | DictDataLoader | None DEFAULT: None

Source code in qadence/ml_tools/callbacks/writer_registry.py

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.")

open(config, iteration=None)

Opens the TensorBoard writer.

PARAMETER	DESCRIPTION
config	Configuration object containing settings for logging. TYPE: TrainConfig
iteration	The iteration step to start logging from. Defaults to None. TYPE: int DEFAULT: None

RETURNS	DESCRIPTION
SummaryWriter	The initialized TensorBoard writer. TYPE: SummaryWriter

Source code in qadence/ml_tools/callbacks/writer_registry.py

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

plot(model, iteration, plotting_functions)

Logs plots of the model using provided plotting functions.

PARAMETER	DESCRIPTION
model	The model to plot. TYPE: Module
iteration	The current iteration number. TYPE: int
plotting_functions	Functions used to generate plots. TYPE: tuple[PlottingFunction, ...]

Source code in qadence/ml_tools/callbacks/writer_registry.py

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"
        )

write(result)

Logs the results of the current iteration to TensorBoard.

PARAMETER	DESCRIPTION
result	The optimization results to log. TYPE: OptimizeResult

Source code in qadence/ml_tools/callbacks/writer_registry.py

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"
        )

get_writer(tracking_tool)

Factory method to get the appropriate writer based on the tracking tool.

PARAMETER	DESCRIPTION
tracking_tool	The experiment tracking tool to use. TYPE: ExperimentTrackingTool

RETURNS	DESCRIPTION
BaseWriter	An instance of the appropriate writer. TYPE: BaseWriter

Source code in qadence/ml_tools/callbacks/writer_registry.py

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}")