Fix saving logic

Author: sovrasov

model_api/python/model_api/models/model.py

@@ -99,14 +99,6 @@ def __init__(self, inference_adapter: InferenceAdapter, configuration: dict = {}
         self.callback_fn = lambda _: None
 
     def get_model(self):
-        model_info = {
-            "model_type": self.__model__,
-        }
-        for name in self.parameters():
-            model_info[name] = getattr(self, name)
-
-        self.inference_adapter.update_model_info(model_info)
-
         return self.inference_adapter.get_model()
 
     @classmethod
@@ -493,4 +485,11 @@ def log_layers_info(self):
             )
 
     def save(self, path: str, weights_path: str = "", version: str = "UNSPECIFIED"):
+        model_info = {
+            "model_type": self.__model__,
+        }
+        for name in self.parameters():
+            model_info[name] = getattr(self, name)
+
+        self.inference_adapter.update_model_info(model_info)
         self.inference_adapter.save_model(path, weights_path, version)

tests/python/funtional/test_save.py

@@ -5,6 +5,8 @@
 
 from pathlib import Path
 
+import onnx
+
 from model_api.models import Model
 from model_api.adapters import ONNXRuntimeAdapter
 from model_api.adapters.utils import load_parameters_from_onnx
@@ -15,14 +17,15 @@ def test_detector_save(tmp_path):
         "ssd_mobilenet_v1_fpn_coco",
         configuration={"mean_values": [0, 0, 0], "confidence_threshold": 0.6},
     )
+    xml_path = str(tmp_path / "a.xml")
+    downloaded.save(xml_path)
+    deserialized = Model.create_model(xml_path)
+
     assert (
-        downloaded.get_model()
+        deserialized.get_model()
         .get_rt_info(["model_info", "embedded_processing"])
         .astype(bool)
     )
-    xml_path = str(tmp_path / "a.xml")
-    downloaded.save(xml_path)
-    deserialized = Model.create_model(xml_path)
     assert type(downloaded) is type(deserialized)
     for attr in downloaded.parameters():
         assert getattr(downloaded, attr) == getattr(deserialized, attr)
@@ -32,14 +35,15 @@ def test_classifier_save(tmp_path):
     downloaded = Model.create_model(
         "efficientnet-b0-pytorch", configuration={"scale_values": [1, 1, 1], "topk": 6}
     )
+    xml_path = str(tmp_path / "a.xml")
+    downloaded.save(xml_path)
+    deserialized = Model.create_model(xml_path)
+
     assert (
-        downloaded.get_model()
+        deserialized.get_model()
         .get_rt_info(["model_info", "embedded_processing"])
         .astype(bool)
     )
-    xml_path = str(tmp_path / "a.xml")
-    downloaded.save(xml_path)
-    deserialized = Model.create_model(xml_path)
     assert type(downloaded) is type(deserialized)
     for attr in downloaded.parameters():
         assert getattr(downloaded, attr) == getattr(deserialized, attr)
@@ -50,14 +54,15 @@ def test_segmentor_save(tmp_path):
         "hrnet-v2-c1-segmentation",
         configuration={"reverse_input_channels": True, "labels": ["first", "second"]},
     )
+    xml_path = str(tmp_path / "a.xml")
+    downloaded.save(xml_path)
+    deserialized = Model.create_model(xml_path)
+
     assert (
-        downloaded.get_model()
+        deserialized.get_model()
         .get_rt_info(["model_info", "embedded_processing"])
         .astype(bool)
     )
-    xml_path = str(tmp_path / "a.xml")
-    downloaded.save(xml_path)
-    deserialized = Model.create_model(xml_path)
     assert type(downloaded) is type(deserialized)
     for attr in downloaded.parameters():
         assert getattr(downloaded, attr) == getattr(deserialized, attr)
@@ -71,17 +76,16 @@ def test_onnx_save(tmp_path, data):
         configuration={"reverse_input_channels": True, "topk": 6},
     )
 
+    onnx_path = str(tmp_path / "a.onnx")
+    cls_model.save(onnx_path)
+    deserialized = Model.create_model(onnx_path)
+
     assert (
-        load_parameters_from_onnx(cls_model.get_model())["model_info"][
+        load_parameters_from_onnx(onnx.load(onnx_path))["model_info"][
             "embedded_processing"
         ]
         == "True"
     )
-
-    onnx_path = str(tmp_path / "a.onnx")
-    cls_model.save(onnx_path)
-
-    deserialized = Model.create_model(onnx_path)
     assert type(cls_model) is type(deserialized)
     for attr in cls_model.parameters():
         assert getattr(cls_model, attr) == getattr(deserialized, attr)