Merge remote-tracking branch 'origin/develop' into dl/torch_experimen…

…tal_statistics

nncf/common/logging/track_progress.py

@@ -114,6 +114,9 @@ def __init__(
                 TimeRemainingColumn(),
             )
         )
+
+        disable = disable or (hasattr(sequence, "__len__") and len(sequence) == 0)
+
         self.progress = Progress(
             *self.columns,
             auto_refresh=auto_refresh,

nncf/openvino/graph/metatypes/openvino_metatypes.py

@@ -673,6 +673,12 @@ class OVAbsMetatype(OVOpMetatype):
     op_names = ["Abs"]
 
 
+@OV_OPERATOR_METATYPES.register()
+class OVIfMetatype(OVOpMetatype):
+    name = "IfOp"
+    op_names = ["If"]
+
+
 @OV_OPERATOR_METATYPES.register()
 class OVGroupNormalizationMetatype(OVOpMetatype):
     name = "GroupNormalizationOp"

nncf/openvino/graph/model_transformer.py

@@ -25,12 +25,14 @@
 from nncf.openvino.graph.node_utils import get_result_node_name
 from nncf.openvino.graph.transformations.commands import OVBiasCorrectionCommand
 from nncf.openvino.graph.transformations.commands import OVBiasInsertionCommand
+from nncf.openvino.graph.transformations.commands import OVExtractIfBodyCommand
 from nncf.openvino.graph.transformations.commands import OVFQNodeRemovingCommand
 from nncf.openvino.graph.transformations.commands import OVInplaceFnInsertionCommand
 from nncf.openvino.graph.transformations.commands import OVModelExtractionCommand
 from nncf.openvino.graph.transformations.commands import OVMultiplyInsertionCommand
 from nncf.openvino.graph.transformations.commands import OVOutputInsertionCommand
 from nncf.openvino.graph.transformations.commands import OVQuantizerInsertionCommand
+from nncf.openvino.graph.transformations.commands import OVUpdateIfBodyCommand
 from nncf.openvino.graph.transformations.commands import OVWeightUpdateCommand
 from nncf.quantization.fake_quantize import FakeQuantizeParameters
 
@@ -52,6 +54,8 @@ def __init__(self, model: TModel):
             (OVOutputInsertionCommand, self._apply_output_insertion_transformations),
             (OVBiasInsertionCommand, self._apply_bias_insertion_transformations),
             (OVMultiplyInsertionCommand, self._apply_multiply_insertion_transformations),
+            (OVUpdateIfBodyCommand, self._apply_update_if_body_transformations),
+            (OVExtractIfBodyCommand, self._apply_extract_if_body_transformation),
         ]
 
     @staticmethod
@@ -526,3 +530,42 @@ def _apply_multiply_insertion_transformations(
                 destination_port.replace_source_output(multiply_node.output(0))
 
         return model
+
+    @staticmethod
+    def _apply_update_if_body_transformations(
+        model: ov.Model, transformations: List[OVUpdateIfBodyCommand]
+    ) -> ov.Model:
+        """
+        Update model body for IF node.
+
+        :param model: Model to update and insert a new subgraph.
+        :param transformations: Transformations with information of If node and an updated subgraph.
+        :return: Original model with an updated subgraph.
+        """
+        name_to_node_mapping = OVModelTransformer._get_name_to_node_mapping(model)
+        for transformation in transformations:
+            subgraph_model = transformation.subgraph_model
+            port_id = transformation.target_point.port_id
+            node_name = transformation.target_point.target_node_name
+            node = name_to_node_mapping[node_name]
+            node.set_function(port_id, subgraph_model)
+        return model
+
+    @staticmethod
+    def _apply_extract_if_body_transformation(
+        model: ov.Model, transformations: List[OVExtractIfBodyCommand]
+    ) -> ov.Model:
+        """
+        Extract a model body from If node.
+
+        :param model: Model from which extracts a subgraph.
+        :param transformations: Transformations with information from which
+        If node and input port extract a model subgraph.
+        :return: Model subgraph.
+        """
+        transformation = transformations[-1]
+        name_to_node_mapping = OVModelTransformer._get_name_to_node_mapping(model)
+        ov_node = name_to_node_mapping[transformation.if_node_name]
+        if transformation.if_body_condition:
+            return ov.Model(ov_node.get_function(0))  # ticket: 121115
+        return ov.Model(ov_node.get_function(1))  # ticket: 121115

nncf/openvino/graph/node_utils.py

@@ -23,7 +23,9 @@
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVAddMetatype
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVConstantMetatype
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVConvertMetatype
+from nncf.openvino.graph.metatypes.openvino_metatypes import OVIfMetatype
 from nncf.openvino.graph.metatypes.openvino_metatypes import OVMatMulMetatype
+from nncf.openvino.graph.metatypes.openvino_metatypes import get_node_metatype
 
 InplaceInsertionFnType = Callable[[ov.Node, int], ov.Node]
 
@@ -49,6 +51,25 @@ def is_node_with_bias(node: NNCFNode, nncf_graph: NNCFGraph) -> bool:
     return bias_constant is not None
 
 
+def get_number_if_op(model: ov.Model) -> int:
+    """
+    Returns number of If operation in a model.
+
+    :param model: Model.
+    :return: True if Model has If operation, False - otherwise.
+    """
+
+    def cnt_if_op(model: ov.Model, cnt: int) -> int:
+        for op in model.get_ops():
+            if get_node_metatype(op) == OVIfMetatype:
+                cnt += 1
+                cnt = cnt_if_op(op.get_function(0), cnt)
+                cnt = cnt_if_op(op.get_function(1), cnt)
+        return cnt
+
+    return cnt_if_op(model, 0)
+
+
 def get_const_value(const_node: ov.Node) -> np.ndarray:
     """
     Returns the constant tensor for the node.

nncf/openvino/graph/transformations/commands.py

@@ -12,6 +12,7 @@
 from typing import List
 
 import numpy as np
+import openvino.runtime as ov
 
 from nncf.common.graph.transformations.commands import Command
 from nncf.common.graph.transformations.commands import TargetPoint
@@ -191,3 +192,40 @@ def __init__(
     def union(self, other: "TransformationCommand") -> "TransformationCommand":
         # Have a look at nncf/torch/graph/transformations/commands/PTInsertionCommand
         raise NotImplementedError()
+
+
+class OVUpdateIfBodyCommand(TransformationCommand):
+    """
+    Updates If node body.
+    """
+
+    def __init__(self, target_point: OVTargetPoint, body_model: ov.Model):
+        """
+        :param target_point: The TargetPoint instance for the change that contains layer's information.
+        :param body_model: A new model to set.
+        """
+        super().__init__(TransformationType.CHANGE, target_point)
+        self.subgraph_model = body_model
+
+    def union(self, other: "TransformationCommand") -> "TransformationCommand":
+        # Have a look at nncf/torch/graph/transformations/commands/PTInsertionCommand
+        raise NotImplementedError()
+
+
+class OVExtractIfBodyCommand(Command):
+    """
+    Extracts If node body.
+    """
+
+    def __init__(self, if_node_name: str, if_body_condition: bool):
+        """
+        :param target_point: The TargetPoint instance for the extraction that contains layer's information.
+        :param if_body_condition: If true extracts then body, else - else body.
+        """
+        super().__init__(TransformationType.EXTRACT)
+        self.if_node_name = if_node_name
+        self.if_body_condition = if_body_condition
+
+    def union(self, other: "TransformationCommand") -> "TransformationCommand":
+        # Have a look at nncf/torch/graph/transformations/commands/PTInsertionCommand
+        raise NotImplementedError()