We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Training process is crashing if OpenVINO is not found in the environment
**Describe the bug**
When OpenVINO is not present in the environment, the training process crashes after a single epoch.

**To Reproduce**
Steps to reproduce the behavior:
1. Ensure OpenVINO is not installed in the environment (even `mo` should not be present)
2. Run `python -c "from testing.test_full import test_train_inference_optimize_segmentation_rad_2d;test_train_inference_optimize_segmentation_rad_2d('cpu')"`
3. See error:
```powershell
[SNIP!]
FileNotFoundError: [WinError 2] The system cannot find the file specified
```
OR:
```bash
[SNIP!]
Best model could not be loaded, error:  [Errno 2] No such file or directory: 'mo'
```

**Expected behavior**
If `mo` or openvino is not found, continue the training process, just without doing any post-training optimization.

**Screenshots**
N.A.

**GaNDLF Version**
<!-- Put the output of the following command:
python -c 'import GANDLF as g;print(g.__version__)'
-->
0.0.16-dev

**Desktop (please complete the following information):**
N.A.

**Additional context**
N.A.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `GANDLF/utils/modelio.py`
Content:
```
1 import os, hashlib, pkg_resources, subprocess
2 import torch
3 
4 from .generic import get_unique_timestamp
5 
6 # these are the base keys for the model dictionary to save
7 model_dict_full = {
8     "epoch": 0,
9     "model_state_dict": None,
10     "optimizer_state_dict": None,
11     "loss": None,
12     "timestamp": None,
13     "timestamp_hash": None,
14     "git_hash": None,
15     "version": None,
16 }
17 
18 model_dict_required = {
19     "model_state_dict": None,
20     "optimizer_state_dict": None,
21 }
22 
23 best_model_path_end = "_best.pth.tar"
24 latest_model_path_end = "_latest.pth.tar"
25 initial_model_path_end = "_initial.pth.tar"
26 
27 
28 def optimize_and_save_model(model, params, path, onnx_export=True):
29     """
30     Perform post-training optimization and save it to a file.
31 
32     Args:
33         model (torch model): Trained torch model.
34         params (dict): The parameter dictionary.
35         path (str): The path to save the model dictionary to.
36         onnx_export (bool): Whether to export to ONNX and OpenVINO.
37     """
38     onnx_export = params["model"].get("onnx_export", onnx_export)
39     # check for incompatible topologies and disable onnx export
40     # customized imagenet_vgg no longer supported for onnx export: https://github.com/pytorch/pytorch/issues/42653
41     if onnx_export:
42         if (params["model"]["architecture"] in ["sdnet", "brain_age"]) or (
43             "imagenet_vgg" in params["model"]["architecture"]
44         ):
45             onnx_export = False
46 
47     if not (onnx_export):
48         if "onnx_print" not in params:
49             print("WARNING: Current model is not supported by ONNX/OpenVINO!")
50             params["onnx_print"] = True
51         return
52     else:
53         try:
54             print("Optimizing best model.")
55             num_channel = params["model"]["num_channels"]
56             model_dimension = params["model"]["dimension"]
57             ov_output_data_type = params["model"].get("data_type", "FP32")
58             input_shape = params["patch_size"]
59             onnx_path = path
60             if not (onnx_path.endswith(".onnx")):
61                 onnx_path = onnx_path.replace("pth.tar", "onnx")
62             if model_dimension == 2:
63                 dummy_input = torch.randn(
64                     (1, num_channel, input_shape[0], input_shape[1])
65                 )
66             else:
67                 dummy_input = torch.randn(
68                     (1, num_channel, input_shape[0], input_shape[1], input_shape[2])
69                 )
70 
71             with torch.no_grad():
72                 torch.onnx.export(
73                     model.to("cpu"),
74                     dummy_input.to("cpu"),
75                     onnx_path,
76                     opset_version=11,
77                     export_params=True,
78                     verbose=True,
79                     input_names=["input"],
80                     output_names=["output"],
81                 )
82 
83             ov_output_dir = os.path.dirname(os.path.abspath(path))
84         except RuntimeWarning:
85             print("WARNING: Cannot export to ONNX model.")
86             return
87 
88         try:
89             if model_dimension == 2:
90                 subprocess.call(
91                     [
92                         "mo",
93                         "--input_model",
94                         "{0}".format(onnx_path),
95                         "--input_shape",
96                         "[1,{0},{1},{2}]".format(
97                             num_channel, input_shape[0], input_shape[1]
98                         ),
99                         "--data_type",
100                         "{0}".format(ov_output_data_type),
101                         "--output_dir",
102                         "{0}".format(ov_output_dir),
103                     ],
104                 )
105             else:
106                 subprocess.call(
107                     [
108                         "mo",
109                         "--input_model",
110                         "{0}".format(onnx_path),
111                         "--input_shape",
112                         "[1,{0},{1},{2},{3}]".format(
113                             num_channel, input_shape[0], input_shape[1], input_shape[2]
114                         ),
115                         "--data_type",
116                         "{0}".format(ov_output_data_type),
117                         "--output_dir",
118                         "{0}".format(ov_output_dir),
119                     ],
120                 )
121         except subprocess.CalledProcessError:
122             print("WARNING: OpenVINO Model Optimizer IR conversion failed.")
123 
124 
125 def save_model(model_dict, model, params, path, onnx_export=True):
126     """
127     Save the model dictionary to a file.
128 
129     Args:
130         model_dict (dict): Model dictionary to save.
131         model (torch model): Trained torch model.
132         params (dict): The parameter dictionary.
133         path (str): The path to save the model dictionary to.
134         onnx_export (bool): Whether to export to ONNX and OpenVINO.
135     """
136     model_dict["timestamp"] = get_unique_timestamp()
137     model_dict["timestamp_hash"] = hashlib.sha256(
138         str(model_dict["timestamp"]).encode("utf-8")
139     ).hexdigest()
140     model_dict["version"] = pkg_resources.require("GANDLF")[0].version
141     try:
142         model_dict["git_hash"] = (
143             subprocess.check_output(["git", "rev-parse", "HEAD"])
144             .decode("ascii")
145             .strip()
146         )
147     except subprocess.CalledProcessError:
148         model_dict["git_hash"] = None
149     torch.save(model_dict, path)
150 
151     # post-training optimization
152     optimize_and_save_model(model, params, path, onnx_export=onnx_export)
153 
154 
155 def load_model(path, device, full_sanity_check=True):
156     """
157     Load a model dictionary from a file.
158 
159     Args:
160         path (str): The path to save the model dictionary to.
161         device (torch.device): The device to run the model on.
162         full_sanity_check (bool): Whether to run full sanity checking on model.
163 
164     Returns:
165         dict: Model dictionary containing model parameters and metadata.
166     """
167     model_dict = torch.load(path, map_location=device)
168 
169     # check if the model dictionary is complete
170     if full_sanity_check:
171         incomplete_keys = [
172             key for key in model_dict_full.keys() if key not in model_dict.keys()
173         ]
174         if len(incomplete_keys) > 0:
175             raise RuntimeWarning(
176                 "Model dictionary is incomplete; the following keys are missing:",
177                 incomplete_keys,
178             )
179 
180     # check if required keys are absent, and if so raise an error
181     incomplete_required_keys = [
182         key for key in model_dict_required.keys() if key not in model_dict.keys()
183     ]
184     if len(incomplete_required_keys) > 0:
185         raise KeyError(
186             "Model dictionary is incomplete; the following keys are missing:",
187             incomplete_required_keys,
188         )
189 
190     return model_dict
191 
192 
193 def load_ov_model(path, device="CPU"):
194     """
195     Load an OpenVINO IR model from an .xml file.
196 
197     Args:
198         path (str): The path to the OpenVINO .xml file.
199         device (str): The device to run inference, can be "CPU", "GPU" or "MULTI:CPU,GPU". Default to be "CPU".
200 
201     Returns:
202         exec_net (OpenVINO executable net): executable OpenVINO model.
203         input_blob (str): Input name.
204         output_blob (str): Output name.
205     """
206 
207     try:
208         from openvino import runtime as ov
209     except ImportError:
210         raise ImportError("OpenVINO inference engine is not configured correctly.")
211 
212     core = ov.Core()
213     if device.lower() == "cuda":
214         device = "GPU"
215 
216     if device == "GPU":
217         core.set_property({"CACHE_DIR": os.path.dirname(os.path.abspath(path))})
218 
219     model = core.read_model(model=path, weights=path.replace("xml", "bin"))
220     compiled_model = core.compile_model(model=model, device_name=device.upper())
221     input_layer = compiled_model.inputs
222     output_layer = compiled_model.outputs
223 
224     return compiled_model, input_layer, output_layer
225
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/GANDLF/utils/modelio.py b/GANDLF/utils/modelio.py
--- a/GANDLF/utils/modelio.py
+++ b/GANDLF/utils/modelio.py
@@ -85,41 +85,50 @@
             print("WARNING: Cannot export to ONNX model.")
             return
 
+        # https://github.com/mlcommons/GaNDLF/issues/605
+        openvino_present = False
         try:
-            if model_dimension == 2:
-                subprocess.call(
-                    [
-                        "mo",
-                        "--input_model",
-                        "{0}".format(onnx_path),
-                        "--input_shape",
-                        "[1,{0},{1},{2}]".format(
-                            num_channel, input_shape[0], input_shape[1]
-                        ),
-                        "--data_type",
-                        "{0}".format(ov_output_data_type),
-                        "--output_dir",
-                        "{0}".format(ov_output_dir),
-                    ],
-                )
-            else:
-                subprocess.call(
-                    [
-                        "mo",
-                        "--input_model",
-                        "{0}".format(onnx_path),
-                        "--input_shape",
-                        "[1,{0},{1},{2},{3}]".format(
-                            num_channel, input_shape[0], input_shape[1], input_shape[2]
-                        ),
-                        "--data_type",
-                        "{0}".format(ov_output_data_type),
-                        "--output_dir",
-                        "{0}".format(ov_output_dir),
-                    ],
-                )
-        except subprocess.CalledProcessError:
-            print("WARNING: OpenVINO Model Optimizer IR conversion failed.")
+            import openvino
+            openvino_present = True
+        except ImportError:
+            print("WARNING: OpenVINO is not present.")
+        
+        if openvino_present:
+            try:
+                if model_dimension == 2:
+                    subprocess.call(
+                        [
+                            "mo",
+                            "--input_model",
+                            "{0}".format(onnx_path),
+                            "--input_shape",
+                            "[1,{0},{1},{2}]".format(
+                                num_channel, input_shape[0], input_shape[1]
+                            ),
+                            "--data_type",
+                            "{0}".format(ov_output_data_type),
+                            "--output_dir",
+                            "{0}".format(ov_output_dir),
+                        ],
+                    )
+                else:
+                    subprocess.call(
+                        [
+                            "mo",
+                            "--input_model",
+                            "{0}".format(onnx_path),
+                            "--input_shape",
+                            "[1,{0},{1},{2},{3}]".format(
+                                num_channel, input_shape[0], input_shape[1], input_shape[2]
+                            ),
+                            "--data_type",
+                            "{0}".format(ov_output_data_type),
+                            "--output_dir",
+                            "{0}".format(ov_output_dir),
+                        ],
+                    )
+            except subprocess.CalledProcessError:
+                print("WARNING: OpenVINO Model Optimizer IR conversion failed.")
 
 
 def save_model(model_dict, model, params, path, onnx_export=True):

{"golden_diff": "diff --git a/GANDLF/utils/modelio.py b/GANDLF/utils/modelio.py\n--- a/GANDLF/utils/modelio.py\n+++ b/GANDLF/utils/modelio.py\n@@ -85,41 +85,50 @@\n             print(\"WARNING: Cannot export to ONNX model.\")\n             return\n \n+        # https://github.com/mlcommons/GaNDLF/issues/605\n+        openvino_present = False\n         try:\n-            if model_dimension == 2:\n-                subprocess.call(\n-                    [\n-                        \"mo\",\n-                        \"--input_model\",\n-                        \"{0}\".format(onnx_path),\n-                        \"--input_shape\",\n-                        \"[1,{0},{1},{2}]\".format(\n-                            num_channel, input_shape[0], input_shape[1]\n-                        ),\n-                        \"--data_type\",\n-                        \"{0}\".format(ov_output_data_type),\n-                        \"--output_dir\",\n-                        \"{0}\".format(ov_output_dir),\n-                    ],\n-                )\n-            else:\n-                subprocess.call(\n-                    [\n-                        \"mo\",\n-                        \"--input_model\",\n-                        \"{0}\".format(onnx_path),\n-                        \"--input_shape\",\n-                        \"[1,{0},{1},{2},{3}]\".format(\n-                            num_channel, input_shape[0], input_shape[1], input_shape[2]\n-                        ),\n-                        \"--data_type\",\n-                        \"{0}\".format(ov_output_data_type),\n-                        \"--output_dir\",\n-                        \"{0}\".format(ov_output_dir),\n-                    ],\n-                )\n-        except subprocess.CalledProcessError:\n-            print(\"WARNING: OpenVINO Model Optimizer IR conversion failed.\")\n+            import openvino\n+            openvino_present = True\n+        except ImportError:\n+            print(\"WARNING: OpenVINO is not present.\")\n+        \n+        if openvino_present:\n+            try:\n+                if model_dimension == 2:\n+                    subprocess.call(\n+                        [\n+                            \"mo\",\n+                            \"--input_model\",\n+                            \"{0}\".format(onnx_path),\n+                            \"--input_shape\",\n+                            \"[1,{0},{1},{2}]\".format(\n+                                num_channel, input_shape[0], input_shape[1]\n+                            ),\n+                            \"--data_type\",\n+                            \"{0}\".format(ov_output_data_type),\n+                            \"--output_dir\",\n+                            \"{0}\".format(ov_output_dir),\n+                        ],\n+                    )\n+                else:\n+                    subprocess.call(\n+                        [\n+                            \"mo\",\n+                            \"--input_model\",\n+                            \"{0}\".format(onnx_path),\n+                            \"--input_shape\",\n+                            \"[1,{0},{1},{2},{3}]\".format(\n+                                num_channel, input_shape[0], input_shape[1], input_shape[2]\n+                            ),\n+                            \"--data_type\",\n+                            \"{0}\".format(ov_output_data_type),\n+                            \"--output_dir\",\n+                            \"{0}\".format(ov_output_dir),\n+                        ],\n+                    )\n+            except subprocess.CalledProcessError:\n+                print(\"WARNING: OpenVINO Model Optimizer IR conversion failed.\")\n \n \n def save_model(model_dict, model, params, path, onnx_export=True):\n", "issue": "Training process is crashing if OpenVINO is not found in the environment\n**Describe the bug**\r\nWhen OpenVINO is not present in the environment, the training process crashes after a single epoch.\r\n\r\n**To Reproduce**\r\nSteps to reproduce the behavior:\r\n1. Ensure OpenVINO is not installed in the environment (even `mo` should not be present)\r\n2. Run `python -c \"from testing.test_full import test_train_inference_optimize_segmentation_rad_2d;test_train_inference_optimize_segmentation_rad_2d('cpu')\"`\r\n3. See error:\r\n```powershell\r\n[SNIP!]\r\nFileNotFoundError: [WinError 2] The system cannot find the file specified\r\n```\r\nOR:\r\n```bash\r\n[SNIP!]\r\nBest model could not be loaded, error:  [Errno 2] No such file or directory: 'mo'\r\n```\r\n\r\n**Expected behavior**\r\nIf `mo` or openvino is not found, continue the training process, just without doing any post-training optimization.\r\n\r\n**Screenshots**\r\nN.A.\r\n\r\n**GaNDLF Version**\r\n<!-- Put the output of the following command:\r\npython -c 'import GANDLF as g;print(g.__version__)'\r\n-->\r\n0.0.16-dev\r\n\r\n**Desktop (please complete the following information):**\r\nN.A.\r\n\r\n**Additional context**\r\nN.A.\n", "before_files": [{"content": "import os, hashlib, pkg_resources, subprocess\nimport torch\n\nfrom .generic import get_unique_timestamp\n\n# these are the base keys for the model dictionary to save\nmodel_dict_full = {\n    \"epoch\": 0,\n    \"model_state_dict\": None,\n    \"optimizer_state_dict\": None,\n    \"loss\": None,\n    \"timestamp\": None,\n    \"timestamp_hash\": None,\n    \"git_hash\": None,\n    \"version\": None,\n}\n\nmodel_dict_required = {\n    \"model_state_dict\": None,\n    \"optimizer_state_dict\": None,\n}\n\nbest_model_path_end = \"_best.pth.tar\"\nlatest_model_path_end = \"_latest.pth.tar\"\ninitial_model_path_end = \"_initial.pth.tar\"\n\n\ndef optimize_and_save_model(model, params, path, onnx_export=True):\n    \"\"\"\n    Perform post-training optimization and save it to a file.\n\n    Args:\n        model (torch model): Trained torch model.\n        params (dict): The parameter dictionary.\n        path (str): The path to save the model dictionary to.\n        onnx_export (bool): Whether to export to ONNX and OpenVINO.\n    \"\"\"\n    onnx_export = params[\"model\"].get(\"onnx_export\", onnx_export)\n    # check for incompatible topologies and disable onnx export\n    # customized imagenet_vgg no longer supported for onnx export: https://github.com/pytorch/pytorch/issues/42653\n    if onnx_export:\n        if (params[\"model\"][\"architecture\"] in [\"sdnet\", \"brain_age\"]) or (\n            \"imagenet_vgg\" in params[\"model\"][\"architecture\"]\n        ):\n            onnx_export = False\n\n    if not (onnx_export):\n        if \"onnx_print\" not in params:\n            print(\"WARNING: Current model is not supported by ONNX/OpenVINO!\")\n            params[\"onnx_print\"] = True\n        return\n    else:\n        try:\n            print(\"Optimizing best model.\")\n            num_channel = params[\"model\"][\"num_channels\"]\n            model_dimension = params[\"model\"][\"dimension\"]\n            ov_output_data_type = params[\"model\"].get(\"data_type\", \"FP32\")\n            input_shape = params[\"patch_size\"]\n            onnx_path = path\n            if not (onnx_path.endswith(\".onnx\")):\n                onnx_path = onnx_path.replace(\"pth.tar\", \"onnx\")\n            if model_dimension == 2:\n                dummy_input = torch.randn(\n                    (1, num_channel, input_shape[0], input_shape[1])\n                )\n            else:\n                dummy_input = torch.randn(\n                    (1, num_channel, input_shape[0], input_shape[1], input_shape[2])\n                )\n\n            with torch.no_grad():\n                torch.onnx.export(\n                    model.to(\"cpu\"),\n                    dummy_input.to(\"cpu\"),\n                    onnx_path,\n                    opset_version=11,\n                    export_params=True,\n                    verbose=True,\n                    input_names=[\"input\"],\n                    output_names=[\"output\"],\n                )\n\n            ov_output_dir = os.path.dirname(os.path.abspath(path))\n        except RuntimeWarning:\n            print(\"WARNING: Cannot export to ONNX model.\")\n            return\n\n        try:\n            if model_dimension == 2:\n                subprocess.call(\n                    [\n                        \"mo\",\n                        \"--input_model\",\n                        \"{0}\".format(onnx_path),\n                        \"--input_shape\",\n                        \"[1,{0},{1},{2}]\".format(\n                            num_channel, input_shape[0], input_shape[1]\n                        ),\n                        \"--data_type\",\n                        \"{0}\".format(ov_output_data_type),\n                        \"--output_dir\",\n                        \"{0}\".format(ov_output_dir),\n                    ],\n                )\n            else:\n                subprocess.call(\n                    [\n                        \"mo\",\n                        \"--input_model\",\n                        \"{0}\".format(onnx_path),\n                        \"--input_shape\",\n                        \"[1,{0},{1},{2},{3}]\".format(\n                            num_channel, input_shape[0], input_shape[1], input_shape[2]\n                        ),\n                        \"--data_type\",\n                        \"{0}\".format(ov_output_data_type),\n                        \"--output_dir\",\n                        \"{0}\".format(ov_output_dir),\n                    ],\n                )\n        except subprocess.CalledProcessError:\n            print(\"WARNING: OpenVINO Model Optimizer IR conversion failed.\")\n\n\ndef save_model(model_dict, model, params, path, onnx_export=True):\n    \"\"\"\n    Save the model dictionary to a file.\n\n    Args:\n        model_dict (dict): Model dictionary to save.\n        model (torch model): Trained torch model.\n        params (dict): The parameter dictionary.\n        path (str): The path to save the model dictionary to.\n        onnx_export (bool): Whether to export to ONNX and OpenVINO.\n    \"\"\"\n    model_dict[\"timestamp\"] = get_unique_timestamp()\n    model_dict[\"timestamp_hash\"] = hashlib.sha256(\n        str(model_dict[\"timestamp\"]).encode(\"utf-8\")\n    ).hexdigest()\n    model_dict[\"version\"] = pkg_resources.require(\"GANDLF\")[0].version\n    try:\n        model_dict[\"git_hash\"] = (\n            subprocess.check_output([\"git\", \"rev-parse\", \"HEAD\"])\n            .decode(\"ascii\")\n            .strip()\n        )\n    except subprocess.CalledProcessError:\n        model_dict[\"git_hash\"] = None\n    torch.save(model_dict, path)\n\n    # post-training optimization\n    optimize_and_save_model(model, params, path, onnx_export=onnx_export)\n\n\ndef load_model(path, device, full_sanity_check=True):\n    \"\"\"\n    Load a model dictionary from a file.\n\n    Args:\n        path (str): The path to save the model dictionary to.\n        device (torch.device): The device to run the model on.\n        full_sanity_check (bool): Whether to run full sanity checking on model.\n\n    Returns:\n        dict: Model dictionary containing model parameters and metadata.\n    \"\"\"\n    model_dict = torch.load(path, map_location=device)\n\n    # check if the model dictionary is complete\n    if full_sanity_check:\n        incomplete_keys = [\n            key for key in model_dict_full.keys() if key not in model_dict.keys()\n        ]\n        if len(incomplete_keys) > 0:\n            raise RuntimeWarning(\n                \"Model dictionary is incomplete; the following keys are missing:\",\n                incomplete_keys,\n            )\n\n    # check if required keys are absent, and if so raise an error\n    incomplete_required_keys = [\n        key for key in model_dict_required.keys() if key not in model_dict.keys()\n    ]\n    if len(incomplete_required_keys) > 0:\n        raise KeyError(\n            \"Model dictionary is incomplete; the following keys are missing:\",\n            incomplete_required_keys,\n        )\n\n    return model_dict\n\n\ndef load_ov_model(path, device=\"CPU\"):\n    \"\"\"\n    Load an OpenVINO IR model from an .xml file.\n\n    Args:\n        path (str): The path to the OpenVINO .xml file.\n        device (str): The device to run inference, can be \"CPU\", \"GPU\" or \"MULTI:CPU,GPU\". Default to be \"CPU\".\n\n    Returns:\n        exec_net (OpenVINO executable net): executable OpenVINO model.\n        input_blob (str): Input name.\n        output_blob (str): Output name.\n    \"\"\"\n\n    try:\n        from openvino import runtime as ov\n    except ImportError:\n        raise ImportError(\"OpenVINO inference engine is not configured correctly.\")\n\n    core = ov.Core()\n    if device.lower() == \"cuda\":\n        device = \"GPU\"\n\n    if device == \"GPU\":\n        core.set_property({\"CACHE_DIR\": os.path.dirname(os.path.abspath(path))})\n\n    model = core.read_model(model=path, weights=path.replace(\"xml\", \"bin\"))\n    compiled_model = core.compile_model(model=model, device_name=device.upper())\n    input_layer = compiled_model.inputs\n    output_layer = compiled_model.outputs\n\n    return compiled_model, input_layer, output_layer\n", "path": "GANDLF/utils/modelio.py"}], "after_files": [{"content": "import os, hashlib, pkg_resources, subprocess\nimport torch\n\nfrom .generic import get_unique_timestamp\n\n# these are the base keys for the model dictionary to save\nmodel_dict_full = {\n    \"epoch\": 0,\n    \"model_state_dict\": None,\n    \"optimizer_state_dict\": None,\n    \"loss\": None,\n    \"timestamp\": None,\n    \"timestamp_hash\": None,\n    \"git_hash\": None,\n    \"version\": None,\n}\n\nmodel_dict_required = {\n    \"model_state_dict\": None,\n    \"optimizer_state_dict\": None,\n}\n\nbest_model_path_end = \"_best.pth.tar\"\nlatest_model_path_end = \"_latest.pth.tar\"\ninitial_model_path_end = \"_initial.pth.tar\"\n\n\ndef optimize_and_save_model(model, params, path, onnx_export=True):\n    \"\"\"\n    Perform post-training optimization and save it to a file.\n\n    Args:\n        model (torch model): Trained torch model.\n        params (dict): The parameter dictionary.\n        path (str): The path to save the model dictionary to.\n        onnx_export (bool): Whether to export to ONNX and OpenVINO.\n    \"\"\"\n    onnx_export = params[\"model\"].get(\"onnx_export\", onnx_export)\n    # check for incompatible topologies and disable onnx export\n    # customized imagenet_vgg no longer supported for onnx export: https://github.com/pytorch/pytorch/issues/42653\n    if onnx_export:\n        if (params[\"model\"][\"architecture\"] in [\"sdnet\", \"brain_age\"]) or (\n            \"imagenet_vgg\" in params[\"model\"][\"architecture\"]\n        ):\n            onnx_export = False\n\n    if not (onnx_export):\n        if \"onnx_print\" not in params:\n            print(\"WARNING: Current model is not supported by ONNX/OpenVINO!\")\n            params[\"onnx_print\"] = True\n        return\n    else:\n        try:\n            print(\"Optimizing best model.\")\n            num_channel = params[\"model\"][\"num_channels\"]\n            model_dimension = params[\"model\"][\"dimension\"]\n            ov_output_data_type = params[\"model\"].get(\"data_type\", \"FP32\")\n            input_shape = params[\"patch_size\"]\n            onnx_path = path\n            if not (onnx_path.endswith(\".onnx\")):\n                onnx_path = onnx_path.replace(\"pth.tar\", \"onnx\")\n            if model_dimension == 2:\n                dummy_input = torch.randn(\n                    (1, num_channel, input_shape[0], input_shape[1])\n                )\n            else:\n                dummy_input = torch.randn(\n                    (1, num_channel, input_shape[0], input_shape[1], input_shape[2])\n                )\n\n            with torch.no_grad():\n                torch.onnx.export(\n                    model.to(\"cpu\"),\n                    dummy_input.to(\"cpu\"),\n                    onnx_path,\n                    opset_version=11,\n                    export_params=True,\n                    verbose=True,\n                    input_names=[\"input\"],\n                    output_names=[\"output\"],\n                )\n\n            ov_output_dir = os.path.dirname(os.path.abspath(path))\n        except RuntimeWarning:\n            print(\"WARNING: Cannot export to ONNX model.\")\n            return\n\n        # https://github.com/mlcommons/GaNDLF/issues/605\n        openvino_present = False\n        try:\n            import openvino\n            openvino_present = True\n        except ImportError:\n            print(\"WARNING: OpenVINO is not present.\")\n        \n        if openvino_present:\n            try:\n                if model_dimension == 2:\n                    subprocess.call(\n                        [\n                            \"mo\",\n                            \"--input_model\",\n                            \"{0}\".format(onnx_path),\n                            \"--input_shape\",\n                            \"[1,{0},{1},{2}]\".format(\n                                num_channel, input_shape[0], input_shape[1]\n                            ),\n                            \"--data_type\",\n                            \"{0}\".format(ov_output_data_type),\n                            \"--output_dir\",\n                            \"{0}\".format(ov_output_dir),\n                        ],\n                    )\n                else:\n                    subprocess.call(\n                        [\n                            \"mo\",\n                            \"--input_model\",\n                            \"{0}\".format(onnx_path),\n                            \"--input_shape\",\n                            \"[1,{0},{1},{2},{3}]\".format(\n                                num_channel, input_shape[0], input_shape[1], input_shape[2]\n                            ),\n                            \"--data_type\",\n                            \"{0}\".format(ov_output_data_type),\n                            \"--output_dir\",\n                            \"{0}\".format(ov_output_dir),\n                        ],\n                    )\n            except subprocess.CalledProcessError:\n                print(\"WARNING: OpenVINO Model Optimizer IR conversion failed.\")\n\n\ndef save_model(model_dict, model, params, path, onnx_export=True):\n    \"\"\"\n    Save the model dictionary to a file.\n\n    Args:\n        model_dict (dict): Model dictionary to save.\n        model (torch model): Trained torch model.\n        params (dict): The parameter dictionary.\n        path (str): The path to save the model dictionary to.\n        onnx_export (bool): Whether to export to ONNX and OpenVINO.\n    \"\"\"\n    model_dict[\"timestamp\"] = get_unique_timestamp()\n    model_dict[\"timestamp_hash\"] = hashlib.sha256(\n        str(model_dict[\"timestamp\"]).encode(\"utf-8\")\n    ).hexdigest()\n    model_dict[\"version\"] = pkg_resources.require(\"GANDLF\")[0].version\n    try:\n        model_dict[\"git_hash\"] = (\n            subprocess.check_output([\"git\", \"rev-parse\", \"HEAD\"])\n            .decode(\"ascii\")\n            .strip()\n        )\n    except subprocess.CalledProcessError:\n        model_dict[\"git_hash\"] = None\n    torch.save(model_dict, path)\n\n    # post-training optimization\n    optimize_and_save_model(model, params, path, onnx_export=onnx_export)\n\n\ndef load_model(path, device, full_sanity_check=True):\n    \"\"\"\n    Load a model dictionary from a file.\n\n    Args:\n        path (str): The path to save the model dictionary to.\n        device (torch.device): The device to run the model on.\n        full_sanity_check (bool): Whether to run full sanity checking on model.\n\n    Returns:\n        dict: Model dictionary containing model parameters and metadata.\n    \"\"\"\n    model_dict = torch.load(path, map_location=device)\n\n    # check if the model dictionary is complete\n    if full_sanity_check:\n        incomplete_keys = [\n            key for key in model_dict_full.keys() if key not in model_dict.keys()\n        ]\n        if len(incomplete_keys) > 0:\n            raise RuntimeWarning(\n                \"Model dictionary is incomplete; the following keys are missing:\",\n                incomplete_keys,\n            )\n\n    # check if required keys are absent, and if so raise an error\n    incomplete_required_keys = [\n        key for key in model_dict_required.keys() if key not in model_dict.keys()\n    ]\n    if len(incomplete_required_keys) > 0:\n        raise KeyError(\n            \"Model dictionary is incomplete; the following keys are missing:\",\n            incomplete_required_keys,\n        )\n\n    return model_dict\n\n\ndef load_ov_model(path, device=\"CPU\"):\n    \"\"\"\n    Load an OpenVINO IR model from an .xml file.\n\n    Args:\n        path (str): The path to the OpenVINO .xml file.\n        device (str): The device to run inference, can be \"CPU\", \"GPU\" or \"MULTI:CPU,GPU\". Default to be \"CPU\".\n\n    Returns:\n        exec_net (OpenVINO executable net): executable OpenVINO model.\n        input_blob (str): Input name.\n        output_blob (str): Output name.\n    \"\"\"\n\n    try:\n        from openvino import runtime as ov\n    except ImportError:\n        raise ImportError(\"OpenVINO inference engine is not configured correctly.\")\n\n    core = ov.Core()\n    if device.lower() == \"cuda\":\n        device = \"GPU\"\n\n    if device == \"GPU\":\n        core.set_property({\"CACHE_DIR\": os.path.dirname(os.path.abspath(path))})\n\n    model = core.read_model(model=path, weights=path.replace(\"xml\", \"bin\"))\n    compiled_model = core.compile_model(model=model, device_name=device.upper())\n    input_layer = compiled_model.inputs\n    output_layer = compiled_model.outputs\n\n    return compiled_model, input_layer, output_layer\n", "path": "GANDLF/utils/modelio.py"}]}