Tests are fixed

openvinotoolkit · Jul 10, 2024 · d8f2c5f · d8f2c5f
1 parent ad383b2
commit d8f2c5f
Show file tree

Hide file tree

Showing 5 changed files with 99 additions and 51 deletions.
diff --git a/nncf/common/graph/patterns/patterns.py b/nncf/common/graph/patterns/patterns.py
@@ -253,16 +253,12 @@ def dump_graph(self, path: str) -> None:
 
 
 def merge_two_types_of_operations(first_op: Dict[str, Any], second_op: Dict[str, Any], label: str) -> Dict[str, Any]:
-    res = dict()
-    any_attr = False
-    for attr in [GraphPattern.TYPE_ATTR, GraphPattern.TYPE_ATTR]:
-        if attr in first_op and attr in second_op:
-            res[attr] = first_op[attr]
-            res[attr].extend(second_op[attr])
-            any_attr = True
-    if not any_attr:
-        raise nncf.InternalError("Incorrect dicts of operations")
-    return res
+    if GraphPattern.TYPE_ATTR in first_op and GraphPattern.TYPE_ATTR in second_op:
+        res = {GraphPattern.TYPE_ATTR: first_op[GraphPattern.TYPE_ATTR]}
+        res[GraphPattern.TYPE_ATTR].extend(second_op[GraphPattern.TYPE_ATTR])
+        res[GraphPattern.LABEL_ATTR] = label
+        return res
+    raise nncf.InternalError("Incorrect dicts of operations")
 
 
 @dataclass

diff --git a/tests/common/quantization/mock_graphs.py b/tests/common/quantization/mock_graphs.py
@@ -21,6 +21,7 @@
 from nncf.common.graph.layer_attributes import BaseLayerAttributes
 from nncf.common.graph.layer_attributes import ConvolutionLayerAttributes
 from nncf.common.graph.layer_attributes import Dtype
+from nncf.common.graph.operator_metatypes import OperatorMetatype
 from nncf.common.graph.operator_metatypes import UnknownMetatype
 from nncf.common.insertion_point_graph import InsertionPointGraph
 from nncf.common.insertion_point_graph import PostHookInsertionPoint
@@ -188,16 +189,26 @@ def get_mock_nncf_node_attrs(op_name=None, scope_str=None, metatype=None, type_=
 
 
 def _add_nodes_with_layer_attrs(
-    nx_graph: nx.DiGraph, node_keys: List[str], layer_attrs: Dict[str, BaseLayerAttributes]
+    nx_graph: nx.DiGraph,
+    node_keys: List[str],
+    layer_attrs: Dict[str, BaseLayerAttributes],
+    metatypes: Dict[str, OperatorMetatype] = None,
 ) -> nx.DiGraph:
     for node_key in node_keys:
-        nx_graph.add_node(node_key, **get_mock_nncf_node_attrs(op_name=node_key))
+        metatype = None
+        if metatypes is not None and node_key in metatypes:
+            metatype = metatypes[node_key]
+        nx_graph.add_node(node_key, **get_mock_nncf_node_attrs(op_name=node_key, metatype=metatype))
+
         if node_key in layer_attrs:
             nx_graph.nodes[node_key][NNCFNode.LAYER_ATTRIBUTES] = layer_attrs[node_key]
+
     return nx_graph
 
 
-def get_mock_model_graph_with_mergeable_pattern() -> NNCFGraph:
+def get_mock_model_graph_with_mergeable_pattern(
+    conv2d_metatype=None, batchnorm_metatype=None, relu_metatype=None
+) -> NNCFGraph:
     mock_nx_graph = nx.DiGraph()
 
     #   (A)
@@ -225,7 +236,12 @@ def get_mock_model_graph_with_mergeable_pattern() -> NNCFGraph:
             padding_values=[0, 0, 0, 0],
         )
     }
-    mock_nx_graph = _add_nodes_with_layer_attrs(mock_nx_graph, node_keys, layer_attrs)
+    metatypes = {
+        "conv2d": conv2d_metatype,
+        "batch_norm": batchnorm_metatype,
+        "relu": relu_metatype,
+    }
+    mock_nx_graph = _add_nodes_with_layer_attrs(mock_nx_graph, node_keys, layer_attrs, metatypes)
 
     mock_nx_graph.add_edges_from(
         [
@@ -238,7 +254,9 @@ def get_mock_model_graph_with_mergeable_pattern() -> NNCFGraph:
     return get_nncf_graph_from_mock_nx_graph(mock_nx_graph)
 
 
-def get_mock_model_graph_with_no_mergeable_pattern() -> NNCFGraph:
+def get_mock_model_graph_with_no_mergeable_pattern(
+    conv2d_metatype=None, batchnorm_metatype=None, relu_metatype=None
+) -> NNCFGraph:
     mock_nx_graph = nx.DiGraph()
 
     #   (A)
@@ -270,7 +288,12 @@ def get_mock_model_graph_with_no_mergeable_pattern() -> NNCFGraph:
             padding_values=[0, 0, 0, 0],
         )
     }
-    mock_nx_graph = _add_nodes_with_layer_attrs(mock_nx_graph, node_keys, layer_attrs)
+    metatypes = {
+        "conv2d": conv2d_metatype,
+        "batch_norm": batchnorm_metatype,
+        "relu": relu_metatype,
+    }
+    mock_nx_graph = _add_nodes_with_layer_attrs(mock_nx_graph, node_keys, layer_attrs, metatypes)
 
     mock_nx_graph.add_edges_from(
         [
@@ -285,7 +308,9 @@ def get_mock_model_graph_with_no_mergeable_pattern() -> NNCFGraph:
     return get_nncf_graph_from_mock_nx_graph(mock_nx_graph)
 
 
-def get_mock_model_graph_with_broken_output_edge_pattern() -> NNCFGraph:
+def get_mock_model_graph_with_broken_output_edge_pattern(
+    conv2d_metatype=None, batchnorm_metatype=None, relu_metatype=None
+) -> NNCFGraph:
     mock_nx_graph = nx.DiGraph()
 
     #   (A)
@@ -314,7 +339,12 @@ def get_mock_model_graph_with_broken_output_edge_pattern() -> NNCFGraph:
             padding_values=[0, 0, 0, 0],
         )
     }
-    mock_nx_graph = _add_nodes_with_layer_attrs(mock_nx_graph, node_keys, layer_attrs)
+    metatypes = {
+        "conv2d": conv2d_metatype,
+        "batch_norm": batchnorm_metatype,
+        "relu": relu_metatype,
+    }
+    mock_nx_graph = _add_nodes_with_layer_attrs(mock_nx_graph, node_keys, layer_attrs, metatypes)
 
     mock_nx_graph.add_edges_from(
         [

diff --git a/.../reference_graphs/quantized/synthetic_model/ShiftScale__normalize__multi_input_branch.dot b/.../reference_graphs/quantized/synthetic_model/ShiftScale__normalize__multi_input_branch.dot
@@ -3,25 +3,23 @@ strict digraph  {
 "1 SymmetricQuantizer/symmetric_quantize_0" [id=1, type=symmetric_quantize];
 "2 ShiftScaleParametrized/clone_0" [id=2, type=clone];
 "3 ShiftScaleParametrized/sub__0" [id=3, type=sub_];
-"4 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/sub__0|OUTPUT]/symmetric_quantize_0" [id=4, type=symmetric_quantize];
-"5 ShiftScaleParametrized/div__0" [id=5, type=div_];
-"6 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0" [id=6, type=symmetric_quantize];
-"7 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_0" [id=7, type=symmetric_quantize];
-"8 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0" [id=8, type=conv2d];
-"9 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_1" [id=9, type=symmetric_quantize];
-"10 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_1" [id=10, type=conv2d];
-"11 /nncf_model_output_0" [id=11, type=nncf_model_output];
-"12 /nncf_model_output_1" [id=12, type=nncf_model_output];
+"4 ShiftScaleParametrized/div__0" [id=4, type=div_];
+"5 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0" [id=5, type=symmetric_quantize];
+"6 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_0" [id=6, type=symmetric_quantize];
+"7 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0" [id=7, type=conv2d];
+"8 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_1" [id=8, type=symmetric_quantize];
+"9 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_1" [id=9, type=conv2d];
+"10 /nncf_model_output_0" [id=10, type=nncf_model_output];
+"11 /nncf_model_output_1" [id=11, type=nncf_model_output];
 "0 /nncf_model_input_0" -> "1 SymmetricQuantizer/symmetric_quantize_0";
 "1 SymmetricQuantizer/symmetric_quantize_0" -> "2 ShiftScaleParametrized/clone_0";
-"1 SymmetricQuantizer/symmetric_quantize_0" -> "10 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_1";
+"1 SymmetricQuantizer/symmetric_quantize_0" -> "9 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_1";
 "2 ShiftScaleParametrized/clone_0" -> "3 ShiftScaleParametrized/sub__0";
-"3 ShiftScaleParametrized/sub__0" -> "4 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/sub__0|OUTPUT]/symmetric_quantize_0";
-"4 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/sub__0|OUTPUT]/symmetric_quantize_0" -> "5 ShiftScaleParametrized/div__0";
-"5 ShiftScaleParametrized/div__0" -> "6 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0";
-"6 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0" -> "8 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0";
-"7 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_0" -> "8 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0";
-"8 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0" -> "11 /nncf_model_output_0";
-"9 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_1" -> "10 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_1";
-"10 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_1" -> "12 /nncf_model_output_1";
+"3 ShiftScaleParametrized/sub__0" -> "4 ShiftScaleParametrized/div__0";
+"4 ShiftScaleParametrized/div__0" -> "5 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0";
+"5 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0" -> "7 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0";
+"6 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_0" -> "7 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0";
+"7 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0" -> "10 /nncf_model_output_0";
+"8 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_1" -> "9 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_1";
+"9 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_1" -> "11 /nncf_model_output_1";
 }
diff --git a/...reference_graphs/quantized/synthetic_model/ShiftScale__normalize__single_input_branch.dot b/...reference_graphs/quantized/synthetic_model/ShiftScale__normalize__single_input_branch.dot
@@ -3,19 +3,17 @@ strict digraph  {
 "1 SymmetricQuantizer/symmetric_quantize_0" [id=1, type=symmetric_quantize];
 "2 ShiftScaleParametrized/clone_0" [id=2, type=clone];
 "3 ShiftScaleParametrized/sub__0" [id=3, type=sub_];
-"4 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/sub__0|OUTPUT]/symmetric_quantize_0" [id=4, type=symmetric_quantize];
-"5 ShiftScaleParametrized/div__0" [id=5, type=div_];
-"6 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0" [id=6, type=symmetric_quantize];
-"7 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_0" [id=7, type=symmetric_quantize];
-"8 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0" [id=8, type=conv2d];
-"9 /nncf_model_output_0" [id=9, type=nncf_model_output];
+"4 ShiftScaleParametrized/div__0" [id=4, type=div_];
+"5 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0" [id=5, type=symmetric_quantize];
+"6 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_0" [id=6, type=symmetric_quantize];
+"7 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0" [id=7, type=conv2d];
+"8 /nncf_model_output_0" [id=8, type=nncf_model_output];
 "0 /nncf_model_input_0" -> "1 SymmetricQuantizer/symmetric_quantize_0";
 "1 SymmetricQuantizer/symmetric_quantize_0" -> "2 ShiftScaleParametrized/clone_0";
 "2 ShiftScaleParametrized/clone_0" -> "3 ShiftScaleParametrized/sub__0";
-"3 ShiftScaleParametrized/sub__0" -> "4 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/sub__0|OUTPUT]/symmetric_quantize_0";
-"4 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/sub__0|OUTPUT]/symmetric_quantize_0" -> "5 ShiftScaleParametrized/div__0";
-"5 ShiftScaleParametrized/div__0" -> "6 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0";
-"6 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0" -> "8 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0";
-"7 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_0" -> "8 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0";
-"8 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0" -> "9 /nncf_model_output_0";
+"3 ShiftScaleParametrized/sub__0" -> "4 ShiftScaleParametrized/div__0";
+"4 ShiftScaleParametrized/div__0" -> "5 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0";
+"5 ShiftScaleParametrized/NNCFNetworkInterface[_nncf]/ModuleDict[external_quantizers]/SymmetricQuantizer[ShiftScaleParametrized/div__0|OUTPUT]/symmetric_quantize_0" -> "7 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0";
+"6 ShiftScaleParametrized/NNCFConv2d[conv]/ModuleDict[pre_ops]/UpdateWeight[0]/SymmetricQuantizer[op]/symmetric_quantize_0" -> "7 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0";
+"7 ShiftScaleParametrized/NNCFConv2d[conv]/conv2d_0" -> "8 /nncf_model_output_0";
 }
diff --git a/tests/torch/test_model_transformer.py b/tests/torch/test_model_transformer.py
@@ -12,6 +12,7 @@
 from collections import Counter
 from copy import deepcopy
 from dataclasses import dataclass
+from functools import partial
 from pathlib import Path
 from typing import List
 
@@ -21,6 +22,7 @@
 import torch.nn.functional as F
 from torch import nn
 
+import nncf.torch.graph.operator_metatypes as om
 from nncf.common.graph.definitions import MODEL_INPUT_OP_NAME
 from nncf.common.graph.definitions import MODEL_OUTPUT_OP_NAME
 from nncf.common.graph.patterns.manager import PatternsManager
@@ -367,9 +369,33 @@ def test_priority(self, target_type, trace_parameters, priority_type):
 
 
 MERGE_PATTERN_TEST_CASES = (
-    [get_mock_model_graph_with_mergeable_pattern, "basic_pattern"],
-    [get_mock_model_graph_with_no_mergeable_pattern, "no_pattern"],
-    [get_mock_model_graph_with_broken_output_edge_pattern, "broken_output_edges_pattern"],
+    [
+        partial(
+            get_mock_model_graph_with_mergeable_pattern,
+            conv2d_metatype=om.PTConv2dMetatype,
+            batchnorm_metatype=om.PTBatchNormMetatype,
+            relu_metatype=om.PTRELUMetatype,
+        ),
+        "basic_pattern",
+    ],
+    [
+        partial(
+            get_mock_model_graph_with_no_mergeable_pattern,
+            conv2d_metatype=om.PTConv2dMetatype,
+            batchnorm_metatype=om.PTBatchNormMetatype,
+            relu_metatype=om.PTRELUMetatype,
+        ),
+        "no_pattern",
+    ],
+    [
+        partial(
+            get_mock_model_graph_with_broken_output_edge_pattern,
+            conv2d_metatype=om.PTConv2dMetatype,
+            batchnorm_metatype=om.PTBatchNormMetatype,
+            relu_metatype=om.PTRELUMetatype,
+        ),
+        "broken_output_edges_pattern",
+    ],
 )