[Feature] Add mlflow experiment tracker

.gitignore

@@ -75,3 +75,6 @@ events.out.tfevents.*
 *.dvi
 
 *.gv
+
+# mlflow
+mlruns/

docs/tutorials/qml/ml_tools.md

@@ -284,3 +284,60 @@ def train(
 
     return model, optimizer
 ```
+
+## MLTools with mlflow
+
+MLTools now offers MLflow support.
+
+
+```python
+from __future__ import annotations
+
+import os
+from itertools import count
+import torch
+from torch.utils.data import DataLoader
+
+from qadence.ml_tools import (
+    TrainConfig,
+    train_with_grad,
+
+)
+from qadence.ml_tools.data import to_dataloader
+from qadence.ml_tools.utils import rand_featureparameters
+from qadence.models import QNN, QuantumModel
+from qadence.types import ExperimentTrackingTool
+from qadence import QuantumCircuit, hea, Z
+
+os.environ['MLFLOW_TRACKING_URI'] = 'sqlite:///mlflow.db'
+os.environ['MLFLOW_EXPERIMENT'] = 'mlflow_demonstration'
+os.environ['MLFLOW_RUN_NAME'] = 'test_0'
+
+# in case you want to track remotely
+#os.environ['MLFLOW_TRACKING_USERNAME'] =
+#s.environ['MLFLOW_TRACKING_PASSWORD'] =
+def dataloader(batch_size: int = 25) -> DataLoader:
+    x = torch.linspace(0, 1, batch_size).reshape(-1, 1)
+    y = torch.cos(x)
+    return to_dataloader(x, y, batch_size=batch_size, infinite=True)
+data = dataloader()
+model = QNN(QuantumCircuit(2, hea(2,1)), observable=Z(0))
+cnt = count()
+criterion = torch.nn.MSELoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+inputs = rand_featureparameters(model, 1)
+
+def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+    next(cnt)
+    out = model.expectation(inputs)
+    loss = criterion(out, torch.rand(1))
+    return loss, {}
+
+config = TrainConfig(
+    folder='mlflow_demonstration', max_iter=10, checkpoint_every=1, write_every=1, tracking_tool=ExperimentTrackingTool.MLFLOW
+)
+train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+
+os.system('mlflow ui --port 5000')
+os.system('mlflow ui --backend-store-uri sqlite:///mlflow.db')
+```

docs/tutorials/qml/mlflow_demonstration.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import os
+import random
+from itertools import count
+
+import numpy as np
+import torch
+from matplotlib import pyplot as plt
+from matplotlib.figure import Figure
+from torch.utils.data import DataLoader
+
+from qadence import QuantumCircuit, Z, hea
+from qadence.constructors import feature_map, hamiltonian_factory
+from qadence.ml_tools import TrainConfig, train_with_grad
+from qadence.ml_tools.data import to_dataloader
+from qadence.ml_tools.utils import rand_featureparameters
+from qadence.models import QNN, QuantumModel
+from qadence.types import ExperimentTrackingTool
+
+os.environ["MLFLOW_TRACKING_URI"] = "sqlite:///mlflow.db"
+os.environ["MLFLOW_EXPERIMENT"] = "mlflow_demonstration"
+os.environ["MLFLOW_RUN_NAME"] = "test_0"
+
+hyperparams = {
+    "seed": 42,
+    "batch_size": 10,
+    "n_qubits": 2,
+    "ansatz_depth": 1,
+    "observable": Z,
+}
+
+np.random.seed(hyperparams["seed"])
+torch.manual_seed(hyperparams["seed"])
+random.seed(hyperparams["seed"])
+
+
+# in case you want to track remotely
+# os.environ['MLFLOW_TRACKING_USERNAME'] =
+# s.environ['MLFLOW_TRACKING_PASSWORD'] =
+def dataloader(batch_size: int = 25) -> DataLoader:
+    x = torch.linspace(0, 1, batch_size).reshape(-1, 1)
+    y = torch.cos(x)
+    return to_dataloader(x, y, batch_size=batch_size, infinite=True)
+
+
+obs = hamiltonian_factory(register=hyperparams["n_qubits"], detuning=hyperparams["observable"])
+
+data = dataloader(hyperparams["batch_size"])
+fm = feature_map(hyperparams["n_qubits"], param="x")
+model = QNN(
+    QuantumCircuit(
+        hyperparams["n_qubits"], fm, hea(hyperparams["n_qubits"], hyperparams["ansatz_depth"])
+    ),
+    observable=obs,
+    inputs=["x"],
+)
+cnt = count()
+criterion = torch.nn.MSELoss()
+optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+inputs = rand_featureparameters(model, 1)
+
+
+def loss_fn(model: QuantumModel, data: torch.Tensor) -> tuple[torch.Tensor, dict]:
+    next(cnt)
+    out = model.expectation(inputs)
+    loss = criterion(out, torch.rand(1))
+    return loss, {}
+
+
+def plot_fn(model: QuantumModel, iteration: int) -> tuple[str, Figure]:
+    descr = f"ufa_prediction_epoch_{iteration}.png"
+    fig, ax = plt.subplots()
+    x = torch.linspace(0, 1, 100).reshape(-1, 1)
+    out = model.expectation(x)
+    ax.plot(x.detach().numpy(), out.detach().numpy())
+    return descr, fig
+
+
+config = TrainConfig(
+    folder="mlflow_demonstration",
+    max_iter=10,
+    checkpoint_every=1,
+    plot_every=2,
+    write_every=1,
+    tracking_tool=ExperimentTrackingTool.MLFLOW,
+    hyperparams=hyperparams,
+    plotting_functions=(plot_fn,),
+)
+train_with_grad(model, data, optimizer, config, loss_fn=loss_fn)
+
+os.system("mlflow ui --port 5000")
+os.system("mlflow ui --backend-store-uri sqlite:///mlflow.db")

pyproject.toml

@@ -74,6 +74,7 @@ horqrux = [
 protocols = ["qadence-protocols"]
 libs = ["qadence-libs"]
 dlprof = ["nvidia-pyindex", "nvidia-dlprof[pytorch]"]
+mlflow = ["mlflow"]
 all = [
   "pulser-core==0.18.0",
   "pulser-simulation==0.18.0",

qadence/ml_tools/config.py

@@ -3,8 +3,17 @@
 import datetime
 import os
 from dataclasses import dataclass
+from logging import getLogger
 from pathlib import Path
 from typing import Callable, Optional
+from uuid import uuid4
+
+from matplotlib.figure import Figure
+from torch.nn import Module
+
+from qadence.types import ExperimentTrackingTool
+
+logger = getLogger(__name__)
 
 
 @dataclass
@@ -29,6 +38,11 @@ class TrainConfig:
     """Write tensorboard logs."""
     checkpoint_every: int = 5000
     """Write model/optimizer checkpoint."""
+    plot_every: Optional[int] = None
+    """Write figures.
+
+    NOTE: currently only works with mlflow.
+    """
     folder: Optional[Path] = None
     """Checkpoint/tensorboard logs folder."""
     create_subfolder_per_run: bool = False
@@ -46,6 +60,14 @@ class TrainConfig:
     """The batch_size to use when passing a list/tuple of torch.Tensors."""
     verbose: bool = True
     """Whether or not to print out metrics values during training."""
+    tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+    """The tracking tool of choice."""
+    hyperparams: Optional[dict] = None
+    """Hyperparameters to track."""
+    plotting_functions: Optional[tuple[Callable[[Module, int], tuple[str, Figure]]]] = None
+    """Functions for in-train plotting."""
+
+    # mlflow_callbacks: list[Callable] = [write_mlflow_figure(), write_x()]
 
     def __post_init__(self) -> None:
         if self.folder:
@@ -60,3 +82,35 @@ def __post_init__(self) -> None:
             self.trainstop_criterion = lambda x: x <= self.max_iter
         if self.validation_criterion is None:
             self.validation_criterion = lambda x: False
+        if self.plot_every and self.tracking_tool != ExperimentTrackingTool.MLFLOW:
+            raise NotImplementedError("In-training plots are only available with mlflow tracking.")
+        if self.plot_every and self.plotting_functions is None:
+            logger.warning("Plots tracking is required, but no plotting functions are provided.")
+
+
+@dataclass
+class MLFlowConfig:
+    """
+    Example:
+
+        export MLFLOW_TRACKING_URI=tracking_uri
+        export MLFLOW_TRACKING_USERNAME=username
+        export MLFLOW_TRACKING_PASSWORD=password
+    """
+
+    MLFLOW_TRACKING_URI: str = os.getenv("MLFLOW_TRACKING_URI", "")
+    MLFLOW_TRACKING_USERNAME: str = os.getenv("MLFLOW_TRACKING_USERNAME", "")
+    MLFLOW_TRACKING_PASSWORD: str = os.getenv("MLFLOW_TRACKING_PASSWORD", "")
+    EXPERIMENT: str = os.getenv("MLFLOW_EXPERIMENT", str(uuid4()))
+    RUN_NAME: str = os.getenv("MLFLOW_RUN_NAME", "test_0")
+
+    def __post_init__(self) -> None:
+        import mlflow
+
+        if self.MLFLOW_TRACKING_USERNAME != "":
+            logger.info(
+                f"Intialized mlflow remote logging for user {self.MLFLOW_TRACKING_USERNAME}."
+            )
+        mlflow.set_tracking_uri(self.MLFLOW_TRACKING_URI)
+        mlflow.set_experiment(self.EXPERIMENT)
+        mlflow.start_run(run_name=self.RUN_NAME, nested=False)

qadence/ml_tools/printing.py

@@ -1,7 +1,13 @@
 from __future__ import annotations
 
+from typing import Any, Callable
+
+from matplotlib.figure import Figure
+from torch.nn import Module
 from torch.utils.tensorboard import SummaryWriter
 
+from qadence.types import ExperimentTrackingTool
+
 
 def print_metrics(loss: float | None, metrics: dict, iteration: int) -> None:
     msg = " ".join(
@@ -19,5 +25,65 @@ def write_tensorboard(
         writer.add_scalar(key, arg, iteration)
 
 
-def log_hyperparams(writer: SummaryWriter, hyperparams: dict, metrics: dict) -> None:
+def log_hyperparams_tensorboard(writer: SummaryWriter, hyperparams: dict, metrics: dict) -> None:
     writer.add_hparams(hyperparams, metrics)
+
+
+def plot_tensorboard(
+    writer: SummaryWriter, iteration: int, plotting_functions: tuple[Callable]
+) -> None:
+    raise NotImplementedError("Plot logging with tensorboard is not implemented")
+
+
+def write_mlflow(writer: Any, loss: float | None, metrics: dict, iteration: int) -> None:
+    writer.log_metrics({"loss": float(loss)}, step=iteration)  # type: ignore
+    writer.log_metrics(metrics, step=iteration)  # logs the single metrics
+
+
+def log_hyperparams_mlflow(writer: Any, hyperparams: dict, metrics: dict) -> None:
+    writer.log_params(hyperparams)  # type: ignore
+
+
+def plot_mlflow(
+    writer: SummaryWriter,
+    model: Module,
+    iteration: int,
+    plotting_functions: tuple[Callable[[Module, int], tuple[str, Figure]]],
+) -> None:
+    for pf in plotting_functions:
+        descr, fig = pf(model, iteration)
+        writer.log_figure(fig, descr)
+
+
+TRACKER_MAPPING = {
+    ExperimentTrackingTool.TENSORBOARD: write_tensorboard,
+    ExperimentTrackingTool.MLFLOW: write_mlflow,
+}
+
+LOGGER_MAPPING = {
+    ExperimentTrackingTool.TENSORBOARD: log_hyperparams_tensorboard,
+    ExperimentTrackingTool.MLFLOW: log_hyperparams_mlflow,
+}
+
+PLOTTER_MAPPING = {
+    ExperimentTrackingTool.TENSORBOARD: plot_tensorboard,
+    ExperimentTrackingTool.MLFLOW: plot_mlflow,
+}
+
+
+def write_tracker(
+    args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return TRACKER_MAPPING[tracking_tool](*args)
+
+
+def log_tracker(
+    args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return LOGGER_MAPPING[tracking_tool](*args)
+
+
+def plot_tracker(
+    args: Any, tracking_tool: ExperimentTrackingTool = ExperimentTrackingTool.TENSORBOARD
+) -> None:
+    return PLOTTER_MAPPING[tracking_tool](*args)  # type: ignore