upgrade openvino version

mlcommons · Jul 28, 2023 · 9c9d9e6 · 9c9d9e6
1 parent 0738c92
commit 9c9d9e6
Show file tree

Hide file tree

Showing 7 changed files with 42 additions and 48 deletions.
diff --git a/.devcontainer/onCreateCommand.sh b/.devcontainer/onCreateCommand.sh
@@ -2,7 +2,7 @@
 
 pip install --upgrade pip
 pip install wheel
-pip install openvino-dev==2022.1.0 # [OPTIONAL] to generate optimized models for inference
+pip install openvino-dev==2023.0.1 # [OPTIONAL] to generate optimized models for inference
 pip install mlcube_docker          # [OPTIONAL] to deploy GaNDLF models as MLCube-compliant Docker containers
 pip install medmnist==2.1.0
 pip install torch==1.13.1+cpu torchvision==0.14.1+cpu torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cpu
diff --git a/.github/workflows/python-test.yml b/.github/workflows/python-test.yml
@@ -72,7 +72,7 @@ jobs:
         sudo apt-get install libvips libvips-tools -y
         python -m pip install --upgrade pip
         python -m pip install wheel
-        python -m pip install openvino-dev==2022.1.0 mlcube_docker
+        python -m pip install openvino-dev==2023.0.1 mlcube_docker
         pip install torch==1.13.1+cpu torchvision==0.14.1+cpu torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cpu
         pip install -e .
     - name: Run generic unit tests

diff --git a/Dockerfile-CPU b/Dockerfile-CPU
@@ -8,7 +8,7 @@ RUN apt-get update && apt-get install -y python3.8 python3-pip libjpeg8-dev zlib
 RUN python3.8 -m pip install --upgrade pip
 # EXPLICITLY install cpu versions of torch/torchvision (not all versions have +cpu modes on PyPI...)
 RUN python3.8 -m pip install torch==1.13.1+cpu torchvision==0.14.1+cpu torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cpu
-RUN python3.8 -m pip install openvino-dev==2022.1.0 opencv-python-headless mlcube_docker
+RUN python3.8 -m pip install openvino-dev==2023.0.1 opencv-python-headless mlcube_docker
 
 # Do some dependency installation separately here to make layer caching more efficient
 COPY ./setup.py ./setup.py

diff --git a/Dockerfile-CUDA11.6 b/Dockerfile-CUDA11.6
@@ -10,7 +10,7 @@ LABEL version=1.0
 RUN apt-get update && apt-get install -y python3.8 python3-pip libjpeg8-dev zlib1g-dev python3-dev libpython3.8-dev libffi-dev libgl1
 RUN python3.8 -m pip install --upgrade pip
 RUN python3.8 -m pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
-RUN python3.8 -m pip install openvino-dev==2022.1.0 opencv-python-headless mlcube_docker
+RUN python3.8 -m pip install openvino-dev==2023.0.1 opencv-python-headless mlcube_docker
 
 # Do some dependency installation separately here to make layer caching more efficient
 COPY ./setup.py ./setup.py

diff --git a/Dockerfile-ROCm b/Dockerfile-ROCm
@@ -8,7 +8,7 @@ LABEL version=1.0
 # The base image contains ROCm, python 3.8 and pytorch already, no need to install those
 RUN python3 -m pip install --upgrade pip
 RUN python3.8 -m pip install torch==1.13.1+rocm5.2 torchvision==0.14.1+rocm5.2 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/rocm5.2
-RUN python3 -m pip install --upgrade pip && python3 -m pip install openvino-dev==2022.1.0 opencv-python-headless mlcube_docker
+RUN python3 -m pip install --upgrade pip && python3 -m pip install openvino-dev==2023.0.1 opencv-python-headless mlcube_docker
 RUN apt-get update && apt-get install -y libgl1
 
 # Do some dependency installation separately here to make layer caching more efficient

diff --git a/GANDLF/utils/modelio.py b/GANDLF/utils/modelio.py
@@ -26,40 +26,48 @@
 initial_model_path_end = "_initial.pth.tar"
 
 
+import os
+import torch
+import subprocess
+
+
 def optimize_and_save_model(model, params, path, onnx_export=True):
     """
     Perform post-training optimization and save it to a file.
 
     Args:
-        model (torch model): Trained torch model.
+        model (torch.nn.Module): Trained torch model.
         params (dict): The parameter dictionary.
         path (str): The path to save the model dictionary to.
         onnx_export (bool): Whether to export to ONNX and OpenVINO.
     """
+    # Check if ONNX export is enabled in the parameter dictionary
     onnx_export = params["model"].get("onnx_export", onnx_export)
-    # check for incompatible topologies and disable onnx export
-    # customized imagenet_vgg no longer supported for onnx export: https://github.com/pytorch/pytorch/issues/42653
+
+    # Check for incompatible topologies and disable ONNX export
+    # Customized imagenet_vgg no longer supported for ONNX export
     if onnx_export:
-        if (params["model"]["architecture"] in ["sdnet", "brain_age"]) or (
-            "imagenet_vgg" in params["model"]["architecture"]
-        ):
+        architecture = params["model"]["architecture"]
+        if architecture in ["sdnet", "brain_age"] or "imagenet_vgg" in architecture:
             onnx_export = False
 
-    if not (onnx_export):
+    if not onnx_export:
+        # Print a warning if ONNX export is disabled and not already warned
         if "onnx_print" not in params:
             print("WARNING: Current model is not supported by ONNX/OpenVINO!")
             params["onnx_print"] = True
         return
     else:
         try:
-            print("Optimizing best model.")
+            print("Optimizing the best model.")
             num_channel = params["model"]["num_channels"]
             model_dimension = params["model"]["dimension"]
             ov_output_data_type = params["model"].get("data_type", "FP32")
             input_shape = params["patch_size"]
             onnx_path = path
-            if not (onnx_path.endswith(".onnx")):
+            if not onnx_path.endswith(".onnx"):
                 onnx_path = onnx_path.replace("pth.tar", "onnx")
+
             if model_dimension == 2:
                 dummy_input = torch.randn(
                     (1, num_channel, input_shape[0], input_shape[1])
@@ -69,6 +77,7 @@ def optimize_and_save_model(model, params, path, onnx_export=True):
                     (1, num_channel, input_shape[0], input_shape[1], input_shape[2])
                 )
 
+            # Export the model to ONNX format
             with torch.no_grad():
                 torch.onnx.export(
                     model.to("cpu"),
@@ -86,52 +95,37 @@ def optimize_and_save_model(model, params, path, onnx_export=True):
             print("WARNING: Cannot export to ONNX model.")
             return
 
-        # https://github.com/mlcommons/GaNDLF/issues/605
+        # Check if OpenVINO is present and try to convert the ONNX model
         openvino_present = False
         try:
-            import openvino
+            import openvino as ov
+            from openvino.tools.mo import convert_model
 
             openvino_present = True
         except ImportError:
             print("WARNING: OpenVINO is not present.")
 
         if openvino_present:
+            xml_path = onnx_path.replace("onnx", "xml")
+            bin_path = onnx_path.replace("onnx", "bin")
             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),
-                        ],
+                    ov_model = convert_model(
+                        onnx_path,
+                        input_shape=(1, num_channel, input_shape[0], input_shape[1]),
                     )
                 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),
-                        ],
+                    ov_model = convert_model(
+                        onnx_path,
+                        input_shape=(
+                            1,
+                            num_channel,
+                            input_shape[0],
+                            input_shape[1],
+                            input_shape[2],
+                        ),
                     )
+                ov.runtime.serialize(ov_model, xml_path=xml_path, bin_path=bin_path)
             except subprocess.CalledProcessError:
                 print("WARNING: OpenVINO Model Optimizer IR conversion failed.")
 

diff --git a/docs/setup.md b/docs/setup.md
@@ -42,7 +42,7 @@ GaNDLF's primary computational foundation is built on PyTorch, and as such it su
 The following dependencies are optional, and are needed for specific features of GaNDLF.
 
 ```bash
-(venv_gandlf) $> pip install openvino-dev==2022.1.0 # [OPTIONAL] to generate post-training optimized models for inference
+(venv_gandlf) $> pip install openvino-dev==2023.0.1 # [OPTIONAL] to generate post-training optimized models for inference
 (venv_gandlf) $> pip install mlcube_docker # [OPTIONAL] to deploy GaNDLF models as MLCube-compliant Docker containers
 ```