From d47e10eec307fa96944877a3645bf8d6e7ded577 Mon Sep 17 00:00:00 2001 From: Aleksander Wennersteen Date: Tue, 23 Jul 2024 12:41:23 +0200 Subject: [PATCH] [Bug] Fix MLFlow import, and setup function that caused tests to fail (#510) Co-authored-by: Aleksander Wennersteen --- qadence/ml_tools/config.py | 8 ++++---- qadence/ml_tools/printing.py | 16 ++++++++++++---- tests/ml_tools/test_logging.py | 19 +++++++++---------- 3 files changed, 25 insertions(+), 18 deletions(-) diff --git a/qadence/ml_tools/config.py b/qadence/ml_tools/config.py index d15ec78b..4388be02 100644 --- a/qadence/ml_tools/config.py +++ b/qadence/ml_tools/config.py @@ -267,14 +267,14 @@ class FeatureMapConfig: num_repeats: int | dict[str, int] = 0 """ - Number of feature map layers repeated in the data reuploadig step. + Number of feature map layers repeated in the data reuploading step. If all are to be repeated the same number of times, then can give a single - `int`. For different number of repeatitions for each feature, provide a dict + `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 repeatitions for that feature. + 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 repeatition. + the data gets uploaded N+1 times. Defaults to no repetition. """ operation: Callable[[Parameter | Basic], AnalogBlock] | Type[RX] | None = None diff --git a/qadence/ml_tools/printing.py b/qadence/ml_tools/printing.py index 047a0a8f..81014e78 100644 --- a/qadence/ml_tools/printing.py +++ b/qadence/ml_tools/printing.py @@ -4,7 +4,6 @@ from typing import Any, Callable, Union from matplotlib.figure import Figure -from mlflow.models import infer_signature from torch import Tensor from torch.nn import Module from torch.utils.data import DataLoader @@ -82,6 +81,7 @@ def plot_mlflow( def log_model_mlflow( writer: Any, model: Module, dataloader: DataLoader | DictDataLoader | None ) -> None: + signature = None if dataloader is not None: xs: InputData xs, *_ = next(iter(dataloader)) @@ -94,9 +94,17 @@ def log_model_mlflow( xs[key] = val.numpy() for key, val in preds.items(): preds[key] = val.detach.numpy() - signature = infer_signature(xs, preds) - else: - signature = None + + try: + from mlflow.models import infer_signature + + signature = infer_signature(xs, preds) + except ImportError: + logger.warning( + "An MLFlow specific function has been called but MLFlow failed to import." + "Please install MLFlow or adjust your code." + ) + writer.pytorch.log_model(model, artifact_path="model", signature=signature) diff --git a/tests/ml_tools/test_logging.py b/tests/ml_tools/test_logging.py index d8669341..640dacd8 100644 --- a/tests/ml_tools/test_logging.py +++ b/tests/ml_tools/test_logging.py @@ -32,7 +32,7 @@ def dataloader(batch_size: int = 25) -> DataLoader: return to_dataloader(x, y, batch_size=batch_size, infinite=True) -def setup(model: Module) -> tuple[Callable, Optimizer]: +def setup_model(model: Module) -> tuple[Callable, Optimizer]: cnt = count() criterion = torch.nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.1) @@ -83,7 +83,7 @@ def clean_artifacts(run: Run) -> None: def test_hyperparams_logging_mlflow(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None: model = BasicQuantumModel - loss_fn, optimizer = setup(model) + loss_fn, optimizer = setup_model(model) hyperparams = {"max_iter": int(10), "lr": 0.1} @@ -113,7 +113,7 @@ def test_hyperparams_logging_mlflow(BasicQuantumModel: QuantumModel, tmp_path: P def test_hyperparams_logging_tensorboard(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None: model = BasicQuantumModel - loss_fn, optimizer = setup(model) + loss_fn, optimizer = setup_model(model) hyperparams = {"max_iter": int(10), "lr": 0.1} @@ -131,8 +131,7 @@ def test_hyperparams_logging_tensorboard(BasicQuantumModel: QuantumModel, tmp_pa def test_model_logging_mlflow_basicQM(BasicQuantumModel: QuantumModel, tmp_path: Path) -> None: model = BasicQuantumModel - - loss_fn, optimizer = setup(model) + loss_fn, optimizer = setup_model(model) config = TrainConfig( folder=tmp_path, @@ -154,7 +153,7 @@ def test_model_logging_mlflow_basicQNN(BasicQNN: QNN, tmp_path: Path) -> None: data = dataloader() model = BasicQNN - loss_fn, optimizer = setup(model) + loss_fn, optimizer = setup_model(model) config = TrainConfig( folder=tmp_path, @@ -176,7 +175,7 @@ def test_model_logging_mlflow_basicAdjQNN(BasicAdjointQNN: QNN, tmp_path: Path) data = dataloader() model = BasicAdjointQNN - loss_fn, optimizer = setup(model) + loss_fn, optimizer = setup_model(model) config = TrainConfig( folder=tmp_path, @@ -199,7 +198,7 @@ def test_model_logging_tensorboard( ) -> None: model = BasicQuantumModel - loss_fn, optimizer = setup(model) + loss_fn, optimizer = setup_model(model) config = TrainConfig( folder=tmp_path, @@ -219,7 +218,7 @@ def test_plotting_mlflow(BasicQNN: QNN, tmp_path: Path) -> None: data = dataloader() model = BasicQNN - loss_fn, optimizer = setup(model) + loss_fn, optimizer = setup_model(model) def plot_model(model: QuantumModel, iteration: int) -> tuple[str, Figure]: descr = f"model_prediction_epoch_{iteration}.png" @@ -267,7 +266,7 @@ def test_plotting_tensorboard(BasicQNN: QNN, tmp_path: Path) -> None: data = dataloader() model = BasicQNN - loss_fn, optimizer = setup(model) + loss_fn, optimizer = setup_model(model) def plot_model(model: QuantumModel, iteration: int) -> tuple[str, Figure]: descr = f"model_prediction_epoch_{iteration}.png"