quadra 2.4.0__py3-none-any.whl → 2.5.0__py3-none-any.whl

quadra/__init__.py

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

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/utils/export.py

@@ -430,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))}

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:

{quadra-2.4.0.dist-info → quadra-2.5.0.dist-info}/METADATA

@@ -1,8 +1,9 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: quadra
-Version: 2.4.0
+Version: 2.5.0
 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.0.dist-info → quadra-2.5.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-quadra/__init__.py,sha256=fv-5hfERt0uLXmjb7dOuh4wKtsXYgvW72hTUW_IvWQo,112
+quadra/__init__.py,sha256=my_GwHSTlCXOqpoWuPRTAF1_rKRUMylmeA4eWoMa96g,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
@@ -36,7 +36,7 @@ quadra/configs/callbacks/all.yaml,sha256=LZx8d0apwv9t0KlKqzFCrYo2NpuNBgcUnHhG_ku
 quadra/configs/callbacks/default.yaml,sha256=ZFPU1bm36hJsxI-85uiJx7TpX7qWkR8dibBKWtES4Yc,1180
 quadra/configs/callbacks/default_anomalib.yaml,sha256=FjjSj6HgMvH18MV4AKR4Ew0pOerPefj9eUGinUkODLE,2256
 quadra/configs/config.yaml,sha256=IULhqUF8Z7Cqr5Xx41EGj8dwtPQWSRiZ5jwzNg0Rjwk,686
-quadra/configs/core/default.yaml,sha256=IfgjKHXuOknq3CKvKKBqMPfiqmZSUOnc80Q2Jkbm7go,239
+quadra/configs/core/default.yaml,sha256=-qXngf5iSRIgPPBNZFu00TeYmE2DC0WuG7x_YPSAG44,264
 quadra/configs/datamodule/base/anomaly.yaml,sha256=CILLAoQHunrT4BN0ynOzizTX29k-7B9vDVIZUYm-cBU,377
 quadra/configs/datamodule/base/classification.yaml,sha256=NYtGk4lmi9Os6bP01AU9xJw3cONRPcGXisdrnoOYpjE,546
 quadra/configs/datamodule/base/multilabel_classification.yaml,sha256=zp7AQx7V6cLJn7zRMKA04zzWKgVq5jTDJMGltqd9OP8,613
@@ -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=dIbhnFPHo2wYoeyE48TeSzGjsf1FowCin3_ASR7BFJc,24621
+quadra/utils/export.py,sha256=fUdcZ2_VKBjuM1yK9nuIEUQ6tNE21SACgBfWhyC5rjw,24651
 quadra/utils/imaging.py,sha256=Cz7sGb_axEmnGcwQJP2djFZpIpGCPFIBGT8NWVV-OOE,866
 quadra/utils/logger.py,sha256=tQJ4xpTAFKx1g-UUm5K1x7zgoP6qoXpcUHQyu0rOr1w,556
 quadra/utils/mlflow.py,sha256=DVso1lxn126hil8i4tTf5WFUPJ8uJNAzNU8OXbXwOzw,3586
@@ -288,13 +288,13 @@ quadra/utils/tests/fixtures/models/classification.py,sha256=5qpyOonqK6W2LCUWEHhm
 quadra/utils/tests/fixtures/models/segmentation.py,sha256=CTNXeEPcFxFq-YcNfQi5DbbytPZwBQaZn5dQq3L41j0,765
 quadra/utils/tests/helpers.py,sha256=9PJlwozUl_lpQW-Ck-tN7sGFcgeieEd3q56aYuwMIlk,2381
 quadra/utils/tests/models.py,sha256=KbAlv_ukxaUYsyVNUO_dM0NyIosx8RpC0EVyF1HvPkM,507
-quadra/utils/utils.py,sha256=3tgj_tFFhKsGNJ9jrmULI9rWxFyhuUe53Y5SBJFkwSM,19124
+quadra/utils/utils.py,sha256=_-iD8MG4g_qrzzMcBrgPWSouQU96jI-5kgYSNYjs4d0,21293
 quadra/utils/validator.py,sha256=wmVXycB90VNyAbKBUVncFCxK4nsYiOWJIY3ISXwxYCY,4632
 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.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.5.0.dist-info/METADATA,sha256=s-yGU-tHTS03MQtY5lXMvPv739qgVmz4ZvdDJs_J4FE,17632
+quadra-2.5.0.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+quadra-2.5.0.dist-info/entry_points.txt,sha256=sRYonBZyx-sAJeWcQNQoVQIU5lm02cnCQt6b15k0WHU,43
+quadra-2.5.0.dist-info/licenses/LICENSE,sha256=8cTbQtcWa02YJoSpMeV_gxj3jpMTkxvl-w3WJ5gV_QE,11342
+quadra-2.5.0.dist-info/RECORD,,

{quadra-2.4.0.dist-info → quadra-2.5.0.dist-info}/WHEEL

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