quadra 2.4.0a0__tar.gz → 2.5.1__tar.gz

{quadra-2.4.0a0 → quadra-2.5.1}/PKG-INFO

@@ -1,8 +1,9 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: quadra
-Version: 2.4.0a0
+Version: 2.5.1
 Summary: Deep Learning experiment orchestration library
 License: Apache-2.0
+License-File: LICENSE
 Keywords: deep learning,experiment,lightning,hydra-core
 Author: Federico Belotti
 Author-email: federico.belotti@orobix.com

{quadra-2.4.0a0 → quadra-2.5.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "quadra"
-version = "2.4.0a0"
+version = "2.5.1"
 description = "Deep Learning experiment orchestration library"
 authors = [
 	"Federico Belotti <federico.belotti@orobix.com>",

{quadra-2.4.0a0 → quadra-2.5.1}/quadra/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.4.0a0"
+__version__ = "2.5.1"
 
 
 def get_version():

{quadra-2.4.0a0 → quadra-2.5.1}/quadra/configs/core/default.yaml

@@ -9,3 +9,4 @@ experiment_path: null
 upload_artifacts: False
 upload_models: ${export.types} # Default behavior in quadra <= 1.5.6
 log_level: info
+mlflow_zip_models: False

{quadra-2.4.0a0 → quadra-2.5.1}/quadra/tasks/anomaly.py

@@ -244,7 +244,7 @@ class AnomalibDetection(Generic[AnomalyDataModuleT], LightningTask[AnomalyDataMo
         ):
             threshold = torch.tensor(100.0)
         else:
-            threshold = self.module.image_metrics.F1Score.threshold
+            threshold = self.module.image_metrics.F1Score.threshold  # type: ignore[union-attr,assignment]
 
         # The output of the prediction is a normalized score so the cumulative histogram is displayed with the
         # normalized scores
@@ -328,6 +328,22 @@ class AnomalibDetection(Generic[AnomalyDataModuleT], LightningTask[AnomalyDataMo
                 else:
                     utils.upload_file_tensorboard(a, tensorboard_logger)
 
+    def execute(self):
+        """Execute the experiment and all the steps."""
+        self.prepare()
+        self.train()
+        # When training in fp16 mixed precision, export function casts model weights from fp32 to fp16,
+        # for this reason, predictions logits could slightly change and predictions could be inconsistent between
+        # test and generated report.
+        # Performing export before test allows to have consistent results in test metrics and generated report.
+        if self.config.export is not None and len(self.config.export.types) > 0:
+            self.export()
+        if self.run_test:
+            self.test()
+        if self.report:
+            self.generate_report()
+        self.finalize()
+
 
 class AnomalibEvaluation(Evaluation[AnomalyDataModule]):
     """Evaluation task for Anomalib.
@@ -445,12 +461,22 @@ class AnomalibEvaluation(Evaluation[AnomalyDataModule]):
         training_threshold = self.model_data[f"{self.training_threshold_type}_threshold"]
         optimal_threshold = self.metadata["threshold"]
 
-        normalized_optimal_threshold = cast(float, normalize_anomaly_score(optimal_threshold, training_threshold))
-
         os.makedirs(os.path.join(self.report_path, "predictions"), exist_ok=True)
         os.makedirs(os.path.join(self.report_path, "heatmaps"), exist_ok=True)
 
         anomaly_scores = self.metadata["anomaly_scores"].cpu().numpy()
+
+        # The reason I have to expand dims and cast the optimal threshold to anomaly_scores dtype is because
+        # of internal roundings performed differently by numpy and python
+        # Particularly the normalized_optimal_threshold computed directly using float values might be higher than the
+        # actual value obtained by the anomaly_scores
+        normalized_optimal_threshold = cast(
+            np.ndarray,
+            normalize_anomaly_score(
+                np.expand_dims(np.array(optimal_threshold, dtype=anomaly_scores.dtype), -1), training_threshold
+            ),
+        ).item()
+
         anomaly_scores = normalize_anomaly_score(anomaly_scores, training_threshold)
 
         if not isinstance(anomaly_scores, np.ndarray):

{quadra-2.4.0a0 → quadra-2.5.1}/quadra/tasks/classification.py

@@ -307,6 +307,14 @@ class Classification(Generic[ClassificationDataModuleT], LightningTask[Classific
         # TODO: What happens if we have 64 precision?
         half_precision = "16" in self.trainer.precision
 
+        example_input: torch.Tensor | None = None
+
+        if hasattr(self.trainer, "datamodule") and hasattr(self.trainer.datamodule, "val_dataset"):
+            # Retrieve a better input to evaluate fp16 performance or efficientnetb0 does not sometimes export properly
+            example_input = self.trainer.datamodule.val_dataset[0][0]
+
+        # Selected rtol and atol are quite high, this is mostly done for efficientnetb0 that seems to be
+        # quite unstable in fp16
         self.model_json, export_paths = export_model(
             config=self.config,
             model=module.model,
@@ -314,6 +322,9 @@ class Classification(Generic[ClassificationDataModuleT], LightningTask[Classific
             half_precision=half_precision,
             input_shapes=input_shapes,
             idx_to_class=idx_to_class,
+            example_inputs=example_input,
+            rtol=0.05,
+            atol=0.01,
         )
 
         if len(export_paths) == 0:
@@ -1136,7 +1147,7 @@ class ClassificationEvaluation(Evaluation[ClassificationDataModuleT]):
             return
 
         if isinstance(self.deployment_model.model.features_extractor, timm.models.resnet.ResNet):
-            target_layers = [cast(BaseNetworkBuilder, self.deployment_model.model).features_extractor.layer4[-1]]
+            target_layers = [cast(BaseNetworkBuilder, self.deployment_model.model).features_extractor.layer4[-1]]  # type: ignore[index]
             self.cam = GradCAM(
                 model=self.deployment_model.model,
                 target_layers=target_layers,

{quadra-2.4.0a0 → quadra-2.5.1}/quadra/utils/anomaly.py

@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+from typing import cast
+
 try:
     from typing import Any, TypeAlias
 except ImportError:
@@ -43,6 +45,39 @@ def normalize_anomaly_score(raw_score: MapOrValue, threshold: float) -> MapOrVal
     else:
         normalized_score = 200.0 - ((raw_score / threshold) * 100.0)
 
+    # Ensures that the normalized scores are consistent with the raw scores
+    # For all the items whose prediction changes after normalization, force the normalized score to be
+    # consistent with the prediction made on the raw score by clipping the score:
+    #   - to 100.0 if the prediction was "anomaly" on the raw score and "good" on the normalized score
+    #   - to 99.99 if the prediction was "good" on the raw score and "anomaly" on the normalized score
+    score = raw_score
+    if isinstance(score, torch.Tensor):
+        score = score.cpu().numpy()
+    # Anomalib classify as anomaly if anomaly_score gte threshold
+    is_anomaly_mask = score >= threshold
+    is_not_anomaly_mask = np.bitwise_not(is_anomaly_mask)
+    if isinstance(normalized_score, torch.Tensor):
+        if normalized_score.dim() == 0:
+            normalized_score = (
+                normalized_score.clamp(min=100.0) if is_anomaly_mask else normalized_score.clamp(max=99.99)
+            )
+        else:
+            normalized_score[is_anomaly_mask] = normalized_score[is_anomaly_mask].clamp(min=100.0)
+            normalized_score[is_not_anomaly_mask] = normalized_score[is_not_anomaly_mask].clamp(max=99.99)
+    elif isinstance(normalized_score, np.ndarray) or np.isscalar(normalized_score):
+        if np.isscalar(normalized_score) or normalized_score.ndim == 0:  # type: ignore[union-attr]
+            normalized_score = (
+                np.clip(normalized_score, a_min=100.0, a_max=None)
+                if is_anomaly_mask
+                else np.clip(normalized_score, a_min=None, a_max=99.99)
+            )
+        else:
+            normalized_score = cast(np.ndarray, normalized_score)
+            normalized_score[is_anomaly_mask] = np.clip(normalized_score[is_anomaly_mask], a_min=100.0, a_max=None)
+            normalized_score[is_not_anomaly_mask] = np.clip(
+                normalized_score[is_not_anomaly_mask], a_min=None, a_max=99.99
+            )
+
     if isinstance(normalized_score, torch.Tensor):
         return torch.clamp(normalized_score, 0.0, 1000.0)
 

{quadra-2.4.0a0 → quadra-2.5.1}/quadra/utils/export.py

@@ -119,6 +119,7 @@ def export_torchscript_model(
     input_shapes: list[Any] | None = None,
     half_precision: bool = False,
     model_name: str = "model.pt",
+    example_inputs: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None,
 ) -> tuple[str, Any] | None:
     """Export a PyTorch model with TorchScript.
 
@@ -128,6 +129,8 @@ def export_torchscript_model(
         output_path: Path to save the model
         half_precision: If True, the model will be exported with half precision
         model_name: Name of the exported model
+        example_inputs: If provided use this to evaluate the model instead of generating random inputs, it's expected to
+            be a list of tensors or a single tensor without batch dimension
 
     Returns:
         If the model is exported successfully, the path to the model and the input shape are returned.
@@ -144,7 +147,32 @@ def export_torchscript_model(
     else:
         model.cpu()
 
-    model_inputs = extract_torch_model_inputs(model, input_shapes, half_precision)
+    batch_size = 1
+    model_inputs: tuple[list[Any] | tuple[Any, ...] | torch.Tensor, list[Any]] | None
+    if example_inputs is not None:
+        if isinstance(example_inputs, Sequence):
+            model_input_tensors = []
+            model_input_shapes = []
+
+            for example_input in example_inputs:
+                new_inp = example_input.to(
+                    device="cuda:0" if half_precision else "cpu",
+                    dtype=torch.float16 if half_precision else torch.float32,
+                )
+                new_inp = new_inp.unsqueeze(0).repeat(batch_size, *(1 for x in new_inp.shape))
+                model_input_tensors.append(new_inp)
+                model_input_shapes.append(new_inp[0].shape)
+
+            model_inputs = (model_input_tensors, [model_input_shapes])
+        else:
+            new_inp = example_inputs.to(
+                device="cuda:0" if half_precision else "cpu",
+                dtype=torch.float16 if half_precision else torch.float32,
+            )
+            new_inp = new_inp.unsqueeze(0).repeat(batch_size, *(1 for x in new_inp.shape))
+            model_inputs = (new_inp, [new_inp[0].shape])
+    else:
+        model_inputs = extract_torch_model_inputs(model, input_shapes, half_precision)
 
     if model_inputs is None:
         return None
@@ -182,6 +210,9 @@ def export_onnx_model(
     input_shapes: list[Any] | None = None,
     half_precision: bool = False,
     model_name: str = "model.onnx",
+    example_inputs: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None,
+    rtol: float = 0.01,
+    atol: float = 5e-3,
 ) -> tuple[str, Any] | None:
     """Export a PyTorch model with ONNX.
 
@@ -192,6 +223,10 @@ def export_onnx_model(
         onnx_config: ONNX export configuration
         half_precision: If True, the model will be exported with half precision
         model_name: Name of the exported model
+        example_inputs: If provided use this to evaluate the model instead of generating random inputs, it's expected to
+            be a list of tensors or a single tensor without batch dimension
+        rtol: Relative tolerance for the ONNX safe export in fp16
+        atol: Absolute tolerance for the ONNX safe export in fp16
     """
     if not ONNX_AVAILABLE:
         log.warning("ONNX is not installed, can not export model in this format.")
@@ -210,9 +245,32 @@ def export_onnx_model(
     else:
         batch_size = 1
 
-    model_inputs = extract_torch_model_inputs(
-        model=model, input_shapes=input_shapes, half_precision=half_precision, batch_size=batch_size
-    )
+    model_inputs: tuple[list[Any] | tuple[Any, ...] | torch.Tensor, list[Any]] | None
+    if example_inputs is not None:
+        if isinstance(example_inputs, Sequence):
+            model_input_tensors = []
+            model_input_shapes = []
+
+            for example_input in example_inputs:
+                new_inp = example_input.to(
+                    device="cuda:0" if half_precision else "cpu",
+                    dtype=torch.float16 if half_precision else torch.float32,
+                )
+                new_inp = new_inp.unsqueeze(0).repeat(batch_size, *(1 for x in new_inp.shape))
+                model_input_tensors.append(new_inp)
+                model_input_shapes.append(new_inp[0].shape)
+
+            model_inputs = (model_input_tensors, [model_input_shapes])
+        else:
+            new_inp = example_inputs.to(
+                device="cuda:0" if half_precision else "cpu",
+                dtype=torch.float16 if half_precision else torch.float32,
+            )
+            new_inp = new_inp.unsqueeze(0).repeat(batch_size, *(1 for x in new_inp.shape))
+            model_inputs = ([new_inp], [new_inp[0].shape])
+    else:
+        model_inputs = extract_torch_model_inputs(model, input_shapes, half_precision)
+
     if model_inputs is None:
         return None
 
@@ -266,6 +324,8 @@ def export_onnx_model(
 
     if isinstance(inp, list):
         inp = tuple(inp)  # onnx doesn't like lists representing tuples of inputs
+    elif isinstance(inp, torch.Tensor):
+        inp = (inp,)
 
     if isinstance(inp, dict):
         raise ValueError("ONNX export does not support model with dict inputs")
@@ -290,6 +350,8 @@ def export_onnx_model(
             onnx_config=onnx_config,
             input_shapes=input_shapes,
             input_names=input_names,
+            rtol=rtol,
+            atol=atol,
         )
 
         if not is_export_ok:
@@ -324,6 +386,8 @@ def _safe_export_half_precision_onnx(
     onnx_config: DictConfig,
     input_shapes: list[Any],
     input_names: list[str],
+    rtol: float = 0.01,
+    atol: float = 5e-3,
 ) -> bool:
     """Check that the exported half precision ONNX model does not contain NaN values. If it does, attempt to export
     the model with a more stable export and overwrite the original model.
@@ -335,6 +399,8 @@ def _safe_export_half_precision_onnx(
         onnx_config: ONNX export configuration
         input_shapes: Input shapes for the model
         input_names: Input names for the model
+        rtol: Relative tolerance to evaluate the model
+        atol: Absolute tolerance to evaluate the model
 
     Returns:
         True if the model is stable or it was possible to export a more stable model, False otherwise.
@@ -364,16 +430,15 @@ def _safe_export_half_precision_onnx(
             export_output = export_onnx_model(
                 model=model,
                 output_path=os.path.dirname(export_model_path),
-                onnx_config=onnx_config,
+                # Force to not simplify fp32 model
+                onnx_config=DictConfig({**onnx_config, "simplify": False}),
                 input_shapes=input_shapes,
                 half_precision=False,
                 model_name=os.path.basename(export_model_path),
             )
-            if export_output is not None:
-                export_model_path, _ = export_output
-            else:
-                log.warning("Failed to export model")
-                return False
+            if export_output is None:
+                # This should not happen
+                raise RuntimeError("Failed to export model")
 
             model_fp32 = onnx.load(export_model_path)
             test_data = {input_names[i]: inp[i].float().cpu().numpy() for i in range(len(inp))}
@@ -381,7 +446,7 @@ def _safe_export_half_precision_onnx(
             with open(os.devnull, "w") as f, contextlib.redirect_stdout(f):
                 # This function prints a lot of information that is not useful for the user
                 model_fp16 = auto_convert_mixed_precision(
-                    model_fp32, test_data, rtol=0.01, atol=5e-3, keep_io_types=False
+                    model_fp32, test_data, rtol=rtol, atol=atol, keep_io_types=False
                 )
             onnx.save(model_fp16, export_model_path)
 
@@ -431,6 +496,9 @@ def export_model(
     input_shapes: list[Any] | None = None,
     idx_to_class: dict[int, str] | None = None,
     pytorch_model_type: Literal["backbone", "model"] = "model",
+    example_inputs: list[Any] | tuple[Any, ...] | torch.Tensor | None = None,
+    rtol: float = 0.01,
+    atol: float = 5e-3,
 ) -> tuple[dict[str, Any], dict[str, str]]:
     """Generate deployment models for the task.
 
@@ -443,6 +511,9 @@ def export_model(
         idx_to_class: Mapping from class index to class name
         pytorch_model_type: Type of the pytorch model config to be exported, if it's backbone on disk we will save the
             config.backbone config, otherwise we will save the config.model
+        example_inputs: If provided use this to evaluate the model instead of generating random inputs
+        rtol: Relative tolerance for the ONNX safe export in fp16
+        atol: Absolute tolerance for the ONNX safe export in fp16
 
     Returns:
         If the model is exported successfully, return a dictionary containing information about the exported model and
@@ -468,6 +539,7 @@ def export_model(
                 input_shapes=input_shapes,
                 output_path=export_folder,
                 half_precision=half_precision,
+                example_inputs=example_inputs,
             )
 
             if out is None:
@@ -495,6 +567,9 @@ def export_model(
                 onnx_config=config.export.onnx,
                 input_shapes=input_shapes,
                 half_precision=half_precision,
+                example_inputs=example_inputs,
+                rtol=rtol,
+                atol=atol,
             )
 
             if out is None:

{quadra-2.4.0a0 → quadra-2.5.1}/quadra/utils/utils.py

@@ -8,10 +8,12 @@ import glob
 import json
 import logging
 import os
+import shutil
 import subprocess
 import sys
 import warnings
 from collections.abc import Iterable, Iterator, Sequence
+from tempfile import TemporaryDirectory
 from typing import Any, cast
 
 import cv2
@@ -299,45 +301,78 @@ def finish(
                     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:
-                    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:
-                        if model_type == "pytorch":
-                            logging.warning("Pytorch format still not supported for mlflow upload")
+                mlflow_zip_models = config.core.get("mlflow_zip_models", False)
+                model_uploaded = False
+                with mlflow.start_run(run_id=mlflow_logger.run_id) as _:
+                    for model_path in deployed_models:
+                        model_type = model_type_from_path(model_path)
+                        model_name = os.path.basename(model_path)
+
+                        if model_type is None:
+                            logging.warning("%s model type not supported", model_path)
                             continue
-
-                        model = quadra_export.import_deployment_model(
-                            model_path,
-                            device=device,
-                            inference_config=config.inference,
-                        )
-
-                        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"] and ONNX_AVAILABLE:
-                            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"
+                        if model_type is not None and model_type in types_to_upload:
+                            if model_type == "pytorch" and not mlflow_zip_models:
+                                logging.warning("Pytorch format still not supported for mlflow upload")
+                                continue
+
+                            if mlflow_zip_models:
+                                with TemporaryDirectory() as temp_dir:
+                                    if model_type == "pytorch" and os.path.isfile(
+                                        os.path.join(export_folder, "model_config.yaml")
+                                    ):
+                                        shutil.copy(model_path, temp_dir)
+                                        shutil.copy(os.path.join(export_folder, "model_config.yaml"), temp_dir)
+                                        shutil.make_archive("assets", "zip", root_dir=temp_dir)
+                                    else:
+                                        shutil.make_archive(
+                                            "assets",
+                                            "zip",
+                                            root_dir=os.path.dirname(model_path),
+                                            base_dir=model_name,
+                                        )
+                                    shutil.move("assets.zip", temp_dir)
+                                    mlflow.pyfunc.log_model(
+                                        artifact_path=model_path,
+                                        loader_module="not.used",
+                                        data_path=os.path.join(temp_dir, "assets.zip"),
+                                        pip_requirements=[""],
                                     )
-                                else:
-                                    model_proto = onnx.load(model.model_path)
-                                    mlflow.onnx.log_model(
-                                        model_proto,
+                                    model_uploaded = True
+                            else:
+                                model = quadra_export.import_deployment_model(
+                                    model_path,
+                                    device=device,
+                                    inference_config=config.inference,
+                                )
+
+                                if model_type in ["torchscript", "pytorch"]:
+                                    signature = infer_signature_model(model.model, inputs)
+                                    mlflow.pytorch.log_model(
+                                        model.model,
                                         artifact_path=model_path,
                                         signature=signature,
                                     )
+                                    model_uploaded = True
+
+                                elif model_type in ["onnx", "simplified_onnx"] and ONNX_AVAILABLE:
+                                    if model.model_path is None:
+                                        logging.warning(
+                                            "Cannot log onnx model on mlflow, \
+                                            BaseEvaluationModel 'model_path' attribute is None"
+                                        )
+                                    else:
+                                        signature = infer_signature_model(model, inputs)
+                                        model_proto = onnx.load(model.model_path)
+                                        mlflow.onnx.log_model(
+                                            model_proto,
+                                            artifact_path=model_path,
+                                            signature=signature,
+                                        )
+                                        model_uploaded = True
+
+                    if model_uploaded:
+                        mlflow.log_artifact(os.path.join(export_folder, "model.json"), export_folder)
 
         if tensorboard_logger is not None:
             config_paths = []
@@ -376,7 +411,7 @@ def model_type_from_path(model_path: str) -> str | None:
       - "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).
+      - "json" if the model has a '.json' extension (JSON).
       - None if model extension is not supported.
 
     Example: