quadra 2.3.2a2__py3-none-any.whl → 2.4.0__py3-none-any.whl

quadra/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.3.2a2"
+__version__ = "2.4.0"
 
 
 def get_version():

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/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.
@@ -381,7 +447,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 +497,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 +512,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 +540,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 +568,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.3.2a2.dist-info → quadra-2.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: quadra
-Version: 2.3.2a2
+Version: 2.4.0
 Summary: Deep Learning experiment orchestration library
 License: Apache-2.0
 Keywords: deep learning,experiment,lightning,hydra-core
@@ -39,7 +39,7 @@ Requires-Dist: onnxsim (==0.4.28) ; extra == "onnx"
 Requires-Dist: opencv_python_headless (>=4.7.0,<4.8.0)
 Requires-Dist: pandas (<2.0)
 Requires-Dist: pillow (>=10,<11)
-Requires-Dist: pydantic (==1.10.10)
+Requires-Dist: pydantic (>=1.10.10)
 Requires-Dist: python_dotenv (>=0.21,<0.22)
 Requires-Dist: pytorch_lightning (>=2.4,<2.5)
 Requires-Dist: rich (>=13.2,<13.3)
@@ -49,11 +49,11 @@ Requires-Dist: seaborn (>=0.12,<0.13)
 Requires-Dist: segmentation_models_pytorch-orobix (==0.3.3.dev1)
 Requires-Dist: tensorboard (>=2.11,<2.12)
 Requires-Dist: timm (==0.9.12)
-Requires-Dist: torch (==2.4.1)
+Requires-Dist: torch (==2.6.0)
 Requires-Dist: torchinfo (>=1.8,<1.9)
 Requires-Dist: torchmetrics (>=0.10,<0.11)
 Requires-Dist: torchsummary (>=1.5,<1.6)
-Requires-Dist: torchvision (>=0.19,<0.20)
+Requires-Dist: torchvision (>=0.21,<0.22)
 Requires-Dist: tripy (>=1.0,<1.1)
 Requires-Dist: typing_extensions (==4.11.0) ; python_version < "3.10"
 Requires-Dist: xxhash (>=3.2,<3.3)

{quadra-2.3.2a2.dist-info → quadra-2.4.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-quadra/__init__.py,sha256=7sypQyykpPeQGER7BmqRc3Z6RChcUEoiVg5zWeB9vn8,114
+quadra/__init__.py,sha256=fv-5hfERt0uLXmjb7dOuh4wKtsXYgvW72hTUW_IvWQo,112
 quadra/callbacks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 quadra/callbacks/anomalib.py,sha256=WLBEGhZA9HoP4Yh9UbbC2GzDOKYTkvU9EY1lkZcV7Fs,11971
 quadra/callbacks/lightning.py,sha256=qvtzDiv8ZUV7K11gKHKWCyo-a9XR_Jm_M-IEicTM1Yo,20242
@@ -250,7 +250,7 @@ quadra/schedulers/warmup.py,sha256=chzzrK7OqqlicBCxiF4CqMYNrWu6nflIbRE-C86Jrw0,4
 quadra/tasks/__init__.py,sha256=tmAfMoH0k3UC7r2pNrgbBa1Pfc3tpLl3IObFF6Z0eRE,820
 quadra/tasks/anomaly.py,sha256=RHeiM1vZF1zsva37iYdiGx_HLgdAp8lXnmUzXja69YU,24638
 quadra/tasks/base.py,sha256=piYlTFtvqH-4s4oEq4GczdAs_gL29UHAJGsOC5Sd3Bc,14187
-quadra/tasks/classification.py,sha256=05l3QM3dsU2yTWhXxNAcJ8sZM0Vbfgey-e5EV6p1TX8,52816
+quadra/tasks/classification.py,sha256=_GQOPMGuOZ_uLA9jFhLEaJkW_Sid_WKHNn9ALXnGNmo,53407
 quadra/tasks/patch.py,sha256=nzo8o-ei7iF1Iarvd8-c08s0Rs_lPvVPDLAbkFMx-Qw,20251
 quadra/tasks/segmentation.py,sha256=9Qy-V0Wvoofl7IrfotnSMgBIXcZd-WfZZtetyqmB0FY,16260
 quadra/tasks/ssl.py,sha256=XsaC9hbhvTA5UfHeRaaCstx9mTYacLRmgoCF5Tj9R5M,20547
@@ -262,7 +262,7 @@ quadra/utils/anomaly.py,sha256=49vFvT5-4SxczsEM2Akcut_M1DDwKlOVdGv36oLTgR0,4067
 quadra/utils/classification.py,sha256=dKFuv4RywWhvhstOnEOnaf-6qcViUK0dTgah9m9mw2Q,24917
 quadra/utils/deprecation.py,sha256=zF_S-yqenaZxRBOudhXts0mX763WjEUWCnHd09TZnwY,852
 quadra/utils/evaluation.py,sha256=oooRJPu1AaHhOwvB1Y6SFjQ645OkgrDzKtUvwWq8oq4,19005
-quadra/utils/export.py,sha256=ghNF8mQw-JjZiVeBJ0y8yIQkx8EG8ssPorn3aaIsgcA,20840
+quadra/utils/export.py,sha256=dIbhnFPHo2wYoeyE48TeSzGjsf1FowCin3_ASR7BFJc,24621
 quadra/utils/imaging.py,sha256=Cz7sGb_axEmnGcwQJP2djFZpIpGCPFIBGT8NWVV-OOE,866
 quadra/utils/logger.py,sha256=tQJ4xpTAFKx1g-UUm5K1x7zgoP6qoXpcUHQyu0rOr1w,556
 quadra/utils/mlflow.py,sha256=DVso1lxn126hil8i4tTf5WFUPJ8uJNAzNU8OXbXwOzw,3586
@@ -293,8 +293,8 @@ quadra/utils/validator.py,sha256=wmVXycB90VNyAbKBUVncFCxK4nsYiOWJIY3ISXwxYCY,463
 quadra/utils/visualization.py,sha256=yYm7lPziUOlybxigZ2qTycNewb67Q80H4hjQGWUh788,16094
 quadra/utils/vit_explainability.py,sha256=Gh6BHaDEzWxOjJp1aqvCxLt9Rb8TXd5uKXOAx7-acUk,13351
 hydra_plugins/quadra_searchpath_plugin.py,sha256=AAn4TzR87zUK7nwSsK-KoqALiPtfQ8FvX3fgZPTGIJ0,1189
-quadra-2.3.2a2.dist-info/LICENSE,sha256=8cTbQtcWa02YJoSpMeV_gxj3jpMTkxvl-w3WJ5gV_QE,11342
-quadra-2.3.2a2.dist-info/METADATA,sha256=NL_uzLkCuKb52aWZ-xMbXc4zBt_egd18lpXNZK4Pqw0,17612
-quadra-2.3.2a2.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-quadra-2.3.2a2.dist-info/entry_points.txt,sha256=sRYonBZyx-sAJeWcQNQoVQIU5lm02cnCQt6b15k0WHU,43
-quadra-2.3.2a2.dist-info/RECORD,,
+quadra-2.4.0.dist-info/LICENSE,sha256=8cTbQtcWa02YJoSpMeV_gxj3jpMTkxvl-w3WJ5gV_QE,11342
+quadra-2.4.0.dist-info/METADATA,sha256=FOt90lNFxRQd84gcN-nFewN-IoclxQ8eDZvhWIeh1Do,17610
+quadra-2.4.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+quadra-2.4.0.dist-info/entry_points.txt,sha256=sRYonBZyx-sAJeWcQNQoVQIU5lm02cnCQt6b15k0WHU,43
+quadra-2.4.0.dist-info/RECORD,,

{quadra-2.3.2a2.dist-info → quadra-2.4.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.1.2
+Generator: poetry-core 2.1.3
 Root-Is-Purelib: true
 Tag: py3-none-any