Added openvino compiler with tensorflow interface

nebullvm/operations/inference_learners/openvino.py

@@ -400,4 +400,4 @@ def run(self, *input_tensors: np.ndarray) -> Tuple[np.ndarray, ...]:
     DeepLearningFramework.PYTORCH: PytorchOpenVinoInferenceLearner,
     DeepLearningFramework.TENSORFLOW: TensorflowOpenVinoInferenceLearner,
     DeepLearningFramework.NUMPY: NumpyOpenVinoInferenceLearner,
-}
+}

nebullvm/operations/optimizations/base.py

@@ -18,6 +18,9 @@
     TFLiteBuildInferenceLearner,
     TensorflowBuildInferenceLearner,
 )
+from nebullvm.operations.inference_learners.openvino import (
+    TensorflowOpenVinoInferenceLearner,
+)
 from nebullvm.operations.measures.measures import MetricDropMeasure
 from nebullvm.operations.measures.utils import (
     compute_relative_difference,
@@ -35,6 +38,7 @@
 )
 from nebullvm.operations.optimizations.compilers.openvino import (
     OpenVINOCompiler,
+    TensorFlowOpenVINOCompiler,
 )
 from nebullvm.operations.optimizations.compilers.pytorch import (
     PytorchBackendCompiler,
@@ -183,7 +187,7 @@ def _optimize(
                             else None,
                             quantization_type=q_type,
                             input_tfms=input_tfms,
-                            onnx_output_path=tmp_dir,
+                            tmp_dir_path=tmp_dir,
                         )
 
                         compiled_model = compiler_op.get_result()
@@ -319,7 +323,10 @@ def get_result(self) -> List:
         DeepLearningFramework.NUMPY: ONNXApacheTVMCompiler,
     },
     ModelCompiler.ONNX_RUNTIME: {DeepLearningFramework.NUMPY: ONNXCompiler},
-    ModelCompiler.OPENVINO: {DeepLearningFramework.NUMPY: OpenVINOCompiler},
+    ModelCompiler.OPENVINO: {
+        DeepLearningFramework.NUMPY: OpenVINOCompiler,
+        DeepLearningFramework.TENSORFLOW: TensorFlowOpenVINOCompiler,
+    },
     ModelCompiler.TFLITE: {
         DeepLearningFramework.TENSORFLOW: TFLiteBackendCompiler
     },
@@ -352,7 +359,8 @@ def get_result(self) -> List:
         DeepLearningFramework.NUMPY: ONNXBuildInferenceLearner
     },
     ModelCompiler.OPENVINO: {
-        DeepLearningFramework.NUMPY: OpenVINOBuildInferenceLearner
+        DeepLearningFramework.NUMPY: OpenVINOBuildInferenceLearner,
+        DeepLearningFramework.TENSORFLOW: TensorflowOpenVinoInferenceLearner,
     },
     ModelCompiler.TFLITE: {
         DeepLearningFramework.TENSORFLOW: TFLiteBuildInferenceLearner

nebullvm/operations/optimizations/compilers/deepsparse.py

@@ -29,7 +29,7 @@ def __init__(self):
     def execute(
         self,
         model: Module,
-        onnx_output_path: str,
+        tmp_dir_path : str,
         model_params: ModelParams,
         quantization_type: QuantizationType = None,
         input_data: DataManager = None,
@@ -39,7 +39,7 @@ def execute(
 
         Args:
             model (torch.nn.Module): The pytorch model.
-            onnx_output_path (str): Path where the converted ONNX model will be
+            tmp_dir_path  (str): Path where the converted ONNX model will be
                 stored.
             model_params (ModelParams): The model parameters.
             quantization_type (QuantizationType): The desired
@@ -60,18 +60,18 @@ def execute(
         )
 
         self.compiled_model = self._compile_model(
-            model, onnx_output_path, input_data, model_params
+            model, tmp_dir_path, input_data, model_params
         )
 
     def _compile_model(
         self,
         model: Union[Module, GraphModule],
-        onnx_output_path: str,
+        tmp_dir_path: str,
         input_data: DataManager,
         model_params: ModelParams,
     ) -> str:
         self.conversion_op.model_name = "model_pruned"
-        onnx_pruned_path = Path(onnx_output_path)
+        onnx_pruned_path = Path(tmp_dir_path)
         self.conversion_op.to(self.device).set_state(
             model, input_data
         ).execute(onnx_pruned_path, model_params)

nebullvm/operations/optimizations/compilers/openvino.py

@@ -14,8 +14,10 @@
 )
 from nebullvm.optional_modules.openvino import (
     Core,
+    Model,
     CompiledModel,
 )
+from nebullvm.optional_modules.tensorflow import tensorflow as tf
 from nebullvm.tools.base import (
     QuantizationType,
     ModelParams,
@@ -25,6 +27,127 @@
 from nebullvm.tools.transformations import MultiStageTransformation
 
 
+
+class TensorFlowOpenVINOCompiler(Compiler):
+    supported_ops = {
+        "cpu": [
+            None,
+            QuantizationType.STATIC,
+            QuantizationType.HALF,
+        ],
+        "gpu": [],
+    }
+
+    def __init__(self):
+        super().__init__()
+
+    def execute(
+        self,
+        model: tf.keras.Model,
+        model_params: ModelParams,
+        tmp_dir_path : str,
+        input_tfms: MultiStageTransformation = None,
+        metric_drop_ths: float = None,
+        quantization_type: QuantizationType = None,
+        input_data: DataManager = None,
+        **kwargs,
+    ):
+        """Compile the input model using TF-OPENVINO library.
+
+        Args:
+            model (tf.keras.Model()): The tensorflow model.
+            model_params (ModelParams): The model parameters.
+            output_file_path (str or Path): Path where storing the output
+            TensorFlow file.
+            input_tfms (MultiStageTransformation, optional): Transformations
+                to be performed to the model's input tensors in order to
+                get the prediction. Default: None.
+            metric_drop_ths (float, optional): Threshold for the accepted drop
+                in terms of precision. Any optimized model with a higher drop
+                will be ignored. Default: None.
+            quantization_type (QuantizationType, optional): The desired
+                quantization algorithm to be used. Default: None.
+            input_data (DataManager): User defined data. Default: None
+        """
+
+        if quantization_type not in self.supported_ops[self.device.value]:
+            self.compiled_model = None
+            return
+
+        if quantization_type is QuantizationType.STATIC and input_data is None:
+            raise ValueError("Input data is required for static quantization.")
+
+        self.logger.info(
+            f"Optimizing with {self.__class__.__name__} and "
+            f"q_type: {quantization_type}."
+        )
+
+        check_quantization(quantization_type, metric_drop_ths)
+        train_input_data = input_data.get_split("train").get_numpy_list(
+            QUANTIZATION_DATA_NUM
+        )
+
+        model.save(tmp_dir_path)
+
+        cmd = [
+            "mo",
+            "--saved_model_dir",
+            str(tmp_dir_path),
+            "--output_dir",
+            str(tmp_dir_path),
+            "--input",
+            ",".join(get_input_names(model)),
+            "--input_shape",
+            ",".join(
+                [
+                    f"{list((model_params.batch_size,) + shape)}"
+                    for shape in model_params.input_sizes
+                ]
+            ),
+        ]
+
+        if quantization_type is QuantizationType.DYNAMIC:
+            return None
+
+        if quantization_type is QuantizationType.HALF:
+            cmd = cmd + ["--data_type", "FP16"]
+
+        process = subprocess.Popen(cmd)
+        process.wait()
+        base_path = Path(model).parent
+        openvino_model_path = base_path / f"{Path(model).stem}.xml"
+        openvino_model_weights = base_path / f"{Path(model).stem}.bin"
+
+        if quantization_type not in [QuantizationType.HALF, None]:
+            openvino_model_path, openvino_model_weights = self._quantize_model(
+                model_topology=str(openvino_model_path),
+                model_weights=str(openvino_model_weights),
+                input_names=get_input_names(model),
+                input_data=train_input_data,
+            )
+
+        self.compiled_model = str(
+            Path(openvino_model_path).parent / Path(openvino_model_path).stem
+        )
+
+    def _compile_model(
+        self,
+        model_name: str,
+        model_weights: str,
+        network_parameters: ModelParams,
+    ) -> CompiledModel:
+        core = Core()
+        model = core.read_model(model=model_name, weights=model_weights)
+
+        dynamic_shape = self._get_dynamic_shape(model, network_parameters)
+
+        if dynamic_shape is not None:
+            model.reshape(dynamic_shape)
+
+        return core.compile_model(model=model, device_name="CPU")
+
+
+
 class OpenVINOCompiler(Compiler):
     supported_ops = {
         "cpu": [
@@ -40,7 +163,7 @@ def __init__(self):
 
     def execute(
         self,
-        model: Union[str, Path],
+        model: tf.keras.Model(),
         model_params: ModelParams,
         input_tfms: MultiStageTransformation = None,
         metric_drop_ths: float = None,

nebullvm/operations/optimizations/optimizers.py

@@ -52,6 +52,8 @@ def _select_compilers_from_hardware(self):
         if tensorflow_is_available():
             compilers.append(ModelCompiler.XLA)
             compilers.append(ModelCompiler.TFLITE)
+        if self.device is Device.CPU and openvino_is_available():
+            compilers.append(ModelCompiler.OPENVINO)
         return compilers
 
 

nebullvm/tools/base.py

@@ -138,4 +138,4 @@ def recursively_dictionarize(element):
     @property
     def input_sizes(self):
         for input_info in self.input_infos:
-            yield input_info.size
+            yield input_info.size