feat: Integrate mlfow upload for onnx models (#101)

* fix: Solve segmentation analysis report slowdows issue

* feat: Add model upload param in core config, support onnx mlflow log

* build: Upload version to 1.5.7

* fix: Refactor tests

* add: Print dataset split count also for segmentation

* feat: Replace default core's upload_models values  with export.types

* docs: Update changelog

CHANGELOG.md

@@ -2,6 +2,16 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+### [1.5.7]
+
+#### Added
+
+- Add upload_models to core config
+
+#### Refactored
+
+- infer_signature_torch_model refactored to infer_signature_model
+
 ### [1.5.6]
 
 #### Added

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "quadra"
-version = "1.5.6"
+version = "1.5.7"
 description = "Deep Learning experiment orchestration library"
 authors = [
   "Federico Belotti <[email protected]>",

quadra/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.5.6"
+__version__ = "1.5.7"
 
 
 def get_version():

quadra/configs/core/default.yaml

@@ -7,4 +7,5 @@ cv2_num_threads: 1
 command: "python "
 experiment_path: null
 upload_artifacts: False
+upload_models: ${export.types} # Default behavior in quadra <= 1.5.6
 log_level: info

quadra/datamodules/base.py

@@ -305,13 +305,18 @@ def save_checkpoint(self) -> None:
             self.data.to_csv(self.dataset_file, index=False)
             log.info("Datamodule checkpoint saved to disk.")
 
-        if "targets" in self.data and not isinstance(self.data["targets"].iloc[0], np.ndarray):
-            # If we find a numpy array target it's very likely one hot encoded, in that case we don't want to print
+        if "targets" in self.data:
+            if isinstance(self.data["targets"].iloc[0], np.ndarray):
+                # If we find a numpy array target it's very likely one hot encoded,
+                # in that case we just print the number of train/val/test samples
+                grouping = ["split"]
+            else:
+                grouping = ["split", "targets"]
             log.info("Dataset Info:")
             split_order = {"train": 0, "val": 1, "test": 2}
             log.info(
                 "\n%s",
-                self.data.groupby(["split", "targets"])
+                self.data.groupby(grouping)
                 .size()
                 .to_frame()
                 .reset_index()

quadra/models/evaluation.py

@@ -306,6 +306,7 @@ def to(self, device: str):
 
     def eval(self):
         """Fake interface to match torch models."""
+        return self
 
     def half(self):
         """Convert model to half precision."""

quadra/utils/mlflow.py

@@ -8,29 +8,28 @@
 except ImportError:
     MLFLOW_AVAILABLE = False
 
-from typing import Any, Sequence, TypeVar
+from typing import Any, Sequence
 
 import torch
 from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import MLFlowLogger
-from torch import nn
 
-NnModuleT = TypeVar("NnModuleT", bound=nn.Module)
+from quadra.models.evaluation import BaseEvaluationModel
 
 
 @torch.inference_mode()
-def infer_signature_torch_model(model: NnModuleT, data: list[Any]) -> ModelSignature | None:
+def infer_signature_model(model: BaseEvaluationModel, data: list[Any]) -> ModelSignature | None:
     """Infer input and output signature for a PyTorch/Torchscript model."""
     model = model.eval()
     model_output = model(*data)
 
     try:
-        output_signature = infer_signature_input_torch(model_output)
+        output_signature = infer_signature_input(model_output)
 
         if len(data) == 1:
-            signature_input = infer_signature_input_torch(data[0])
+            signature_input = infer_signature_input(data[0])
         else:
-            signature_input = infer_signature_input_torch(data)
+            signature_input = infer_signature_input(data)
     except ValueError:
         # TODO: Solve circular import as it is not possible to import get_logger right now
         # log.warning("Unable to infer signature for model output type %s", type(model_output))
@@ -39,7 +38,7 @@ def infer_signature_torch_model(model: NnModuleT, data: list[Any]) -> ModelSigna
     return infer_signature(signature_input, output_signature)
 
 
-def infer_signature_input_torch(input_tensor: Any) -> Any:
+def infer_signature_input(input_tensor: Any) -> Any:
     """Recursively infer the signature input format to pass to mlflow.models.infer_signature.
 
     Raises:
@@ -58,7 +57,7 @@ def infer_signature_input_torch(input_tensor: Any) -> Any:
                 # Nested dicts are not supported
                 raise ValueError("Nested dicts are not supported")
 
-            signature[f"output_{i}"] = infer_signature_input_torch(x)
+            signature[f"output_{i}"] = infer_signature_input(x)
     elif isinstance(input_tensor, torch.Tensor):
         signature = input_tensor.cpu().numpy()
     elif isinstance(input_tensor, dict):
@@ -71,7 +70,7 @@ def infer_signature_input_torch(input_tensor: Any) -> Any:
                 # Nested signature is currently not supported by mlflow
                 raise ValueError("Nested sequences are not supported")
 
-            signature[k] = infer_signature_input_torch(v)
+            signature[k] = infer_signature_input(v)
     else:
         raise ValueError(f"Unable to infer signature for model output type {type(input_tensor)}")
 

quadra/utils/utils.py

@@ -14,6 +14,7 @@
 import dotenv
 import mlflow
 import numpy as np
+import onnx
 import pytorch_lightning as pl
 import rich.syntax
 import rich.tree
@@ -28,7 +29,7 @@
 import quadra
 import quadra.utils.export as quadra_export
 from quadra.callbacks.mlflow import get_mlflow_logger
-from quadra.utils.mlflow import infer_signature_torch_model
+from quadra.utils.mlflow import infer_signature_model
 
 IMAGE_EXTENSIONS: List[str] = [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".pbm", ".pgm", ".ppm", ".pxm", ".pnm"]
 
@@ -291,20 +292,40 @@ def finish(
                     List[Any],
                     quadra_export.generate_torch_inputs(input_size, device=device, half_precision=half_precision),
                 )
+                types_to_upload = config.core.get("upload_models")
                 for model_path in deployed_models:
-                    if model_path.endswith(".pt"):
+                    model_type = model_type_from_path(model_path)
+                    if model_type is None:
+                        logging.warning("%s model type not supported", model_path)
+                        continue
+                    if model_type is not None and model_type in types_to_upload:
                         model = quadra_export.import_deployment_model(
                             model_path, device=device, inference_config=config.inference
-                        ).model
-
-                        signature = infer_signature_torch_model(model, inputs)
-
-                        with mlflow.start_run(run_id=mlflow_logger.run_id) as _:
-                            mlflow.pytorch.log_model(
-                                model,
-                                artifact_path=model_path,
-                                signature=signature,
-                            )
+                        )
+
+                        if model_type in ["torchscript", "pytorch"]:
+                            signature = infer_signature_model(model.model, inputs)
+                            with mlflow.start_run(run_id=mlflow_logger.run_id) as _:
+                                mlflow.pytorch.log_model(
+                                    model.model,
+                                    artifact_path=model_path,
+                                    signature=signature,
+                                )
+                        elif model_type in ["onnx", "simplified_onnx"]:
+                            signature = infer_signature_model(model, inputs)
+                            with mlflow.start_run(run_id=mlflow_logger.run_id) as _:
+                                if model.model_path is None:
+                                    logging.warning(
+                                        "Cannot log onnx model on mlflow, \
+                                        BaseEvaluationModel 'model_path' attribute is None"
+                                    )
+                                else:
+                                    model_proto = onnx.load(model.model_path)
+                                    mlflow.onnx.log_model(
+                                        model_proto,
+                                        artifact_path=model_path,
+                                        signature=signature,
+                                    )
 
         if tensorboard_logger is not None:
             config_paths = []
@@ -331,6 +352,41 @@ def load_envs(env_file: Optional[str] = None) -> None:
     dotenv.load_dotenv(dotenv_path=env_file, override=True)
 
 
+def model_type_from_path(model_path: str) -> Optional[str]:
+    """Determine the type of the machine learning model based on its file extension.
+
+    Parameters:
+    - model_path (str): The file path of the machine learning model.
+
+    Returns:
+    - str: The type of the model, which can be one of the following:
+      - "torchscript" if the model has a '.pt' extension (TorchScript).
+      - "pytorch" if the model has a '.pth' extension (PyTorch).
+      - "simplified_onnx" if the model file ends with 'simplified.onnx' (Simplified ONNX).
+      - "onnx" if the model has a '.onnx' extension (ONNX).
+      - "json" id the model has a '.json' extension (JSON).
+      - None if model extension is not supported.
+
+    Example:
+    ```python
+    model_path = "path/to/your/model.onnx"
+    model_type = model_type_from_path(model_path)
+    print(f"The model type is: {model_type}")
+    ```
+    """
+    if model_path.endswith(".pt"):
+        return "torchscript"
+    if model_path.endswith(".pth"):
+        return "pytorch"
+    if model_path.endswith("simplified.onnx"):
+        return "simplified_onnx"
+    if model_path.endswith(".onnx"):
+        return "onnx"
+    if model_path.endswith(".json"):
+        return "json"
+    return None
+
+
 def setup_opencv() -> None:
     """Setup OpenCV to use only one thread and not use OpenCL."""
     cv2.setNumThreads(1)

quadra/utils/visualization.py

@@ -69,6 +69,8 @@ def create_grid_figure(
         fig_size (Tuple[int, int], optional): Figure size. Defaults to (12, 8).
         bounds (Optional[List[Tuple[float, float]]], optional): Bounds for the images. Defaults to None.
     """
+    default_plt_backend = plt.get_backend()
+    plt.switch_backend("Agg")
     _, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=fig_size, squeeze=False)
     for i, row in enumerate(images):
         for j, image in enumerate(row):
@@ -84,6 +86,7 @@ def create_grid_figure(
     plt.tight_layout()
     plt.savefig(file_path, bbox_inches="tight", dpi=300, facecolor="white", transparent=False)
     plt.close()
+    plt.switch_backend(default_plt_backend)
 
 
 def create_visualization_dataset(dataset: torch.utils.data.Dataset):

tests/utilities/test_mlflow_export.py

@@ -18,7 +18,7 @@
 
 from quadra.models.base import ModelSignatureWrapper
 from quadra.utils.export import generate_torch_inputs
-from quadra.utils.mlflow import infer_signature_input_torch, infer_signature_torch_model
+from quadra.utils.mlflow import infer_signature_input, infer_signature_model
 from quadra.utils.tests.models import DoubleInputModel, SingleInputModel
 
 
@@ -57,7 +57,7 @@ def test_single_tensor_signature():
 
     inputs = generate_torch_inputs(model.input_shapes, device="cpu")
 
-    signature = infer_signature_torch_model(model, inputs)
+    signature = infer_signature_model(model, inputs)
 
     expected_input_signature = [TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype)]
 
@@ -76,7 +76,7 @@ def test_multiple_tensor_signature():
     assert model.input_shapes == [(3, 224, 224), (3, 448, 448)]
     inputs = generate_torch_inputs(model.input_shapes, device="cpu")
 
-    signature = infer_signature_torch_model(model, inputs)
+    signature = infer_signature_model(model, inputs)
 
     expected_input_signature = [
         TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="output_0"),
@@ -98,7 +98,7 @@ def test_dict_signature():
 
     inputs = generate_torch_inputs(model.input_shapes, device="cpu")
 
-    signature = infer_signature_torch_model(model, inputs)
+    signature = infer_signature_model(model, inputs)
 
     expected_input_signature = [TensorSpec(shape=(-1, *x.shape[1:]), type=x.numpy().dtype, name="x")]
 
@@ -126,9 +126,9 @@ def test_nested_tuple_signature():
 
     # Nested structures are not supported
     with raises(ValueError):
-        infer_signature_input_torch(inputs)
+        infer_signature_input(inputs)
 
-    signature = infer_signature(infer_signature_input_torch(outputs), infer_signature_input_torch(outputs))
+    signature = infer_signature(infer_signature_input(outputs), infer_signature_input(outputs))
 
     assert check_signature_equality(signature.outputs.inputs, expected_output_signature)
 
@@ -151,9 +151,9 @@ def test_nested_list_signature():
 
     # Nested structures are not supported
     with raises(ValueError):
-        infer_signature_input_torch(inputs)
+        infer_signature_input(inputs)
 
-    signature = infer_signature(infer_signature_input_torch(outputs), infer_signature_input_torch(outputs))
+    signature = infer_signature(infer_signature_input(outputs), infer_signature_input(outputs))
 
     assert check_signature_equality(signature.outputs.inputs, expected_output_signature)
 
@@ -170,10 +170,10 @@ def test_nested_dicts_signature():
 
     # Nested structures are not supported
     with raises(ValueError):
-        infer_signature_input_torch(inputs)
+        infer_signature_input(inputs)
 
     with raises(ValueError):
-        infer_signature_input_torch(outputs)
+        infer_signature_input(outputs)
 
 
 @torch.inference_mode()
@@ -189,7 +189,7 @@ def test_tuple_of_dicts_signature():
 
     # Nested structures are not supported
     with raises(ValueError):
-        infer_signature_input_torch(inputs)
+        infer_signature_input(inputs)
 
     with raises(ValueError):
-        infer_signature_input_torch(outputs)
+        infer_signature_input(outputs)