Use both nncf_graph and inference_nncf_graph

Author: daniil-lyakhov

nncf/quantization/algorithms/min_max/backend.py

@@ -297,10 +297,12 @@ def get_ignored_names_by_layer_attributes(nncf_graph: NNCFGraph) -> Set[str]:
         """
 
     @abstractmethod
-    def get_weight_nodes(self, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
+    def get_weight_nodes(self, nncf_graph: NNCFGraph, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
         """
         Returns nodes that have weights.
 
+        :param nncf_graph: Instance of original NNCFGraph,
+            which contains shape of and constant subgraphs.
         :param inference_nncf_graph: Instance of inference NNCFGraph,
             which does not contain shape of and constant subgraphs.
         :return: All nodes with weights.

nncf/quantization/algorithms/min_max/onnx_backend.py

@@ -217,7 +217,7 @@ def get_ignored_metatypes(model_type: ModelType, device: TargetDevice) -> List[O
     def get_ignored_names_by_layer_attributes(nncf_graph: NNCFGraph) -> Set[str]:
         return set()
 
-    def get_weight_nodes(self, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
+    def get_weight_nodes(self, nncf_grpah: NNCFGraph, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
         return [node for node in inference_nncf_graph.get_all_nodes() if node.layer_attributes.has_weight()]
 
     @staticmethod

nncf/quantization/algorithms/min_max/openvino_backend.py

@@ -215,7 +215,7 @@ def get_ignored_names_by_layer_attributes(nncf_graph: NNCFGraph) -> Set[str]:
                 ignored_names.add(node.node_name)
         return ignored_names
 
-    def get_weight_nodes(self, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
+    def get_weight_nodes(self, nncf_grpah: NNCFGraph, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
         return [
             node
             for node in inference_nncf_graph.get_all_nodes()

nncf/quantization/algorithms/min_max/torch_backend.py

@@ -37,6 +37,7 @@
 from nncf.torch.graph.graph import PTNNCFGraph
 from nncf.torch.graph.graph import PTTargetPoint
 from nncf.torch.graph.operator_metatypes import ELEMENTWISE_OPERATIONS
+from nncf.torch.graph.operator_metatypes import MATMUL_METATYPES
 from nncf.torch.graph.transformations.command_creation import create_quantizer_insertion_command
 from nncf.torch.graph.transformations.command_creation import create_shared_quantizer_insertion_command
 from nncf.torch.graph.transformations.commands import PTInsertionCommand
@@ -51,7 +52,8 @@
 from nncf.torch.quantization.layers import BaseQuantizer
 from nncf.torch.quantization.layers import PTQuantizerSpec
 from nncf.torch.quantization.layers import get_scale_shape
-from nncf.torch.utils import get_weight_nodes_in_inference_graph
+from nncf.torch.utils import get_weight_nodes
+from nncf.torch.utils import is_matmul_with_constant
 
 
 class PTMinMaxAlgoBackend(MinMaxAlgoBackend):
@@ -66,7 +68,7 @@ def preserved_metatypes(self) -> List[OperatorMetatype]:
 
     @property
     def mat_mul_metatypes(self) -> List[OperatorMetatype]:
-        return [om.PTLinearMetatype, om.PTMatMulMetatype, om.PTAddmmMetatype]
+        return MATMUL_METATYPES
 
     @property
     def post_processing_metatypes(self) -> List[OperatorMetatype]:
@@ -341,8 +343,8 @@ def get_ignored_metatypes(model_type: ModelType, device: TargetDevice) -> List[O
     def get_ignored_names_by_layer_attributes(nncf_graph: NNCFGraph) -> Set[str]:
         return set()
 
-    def get_weight_nodes(self, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
-        return get_weight_nodes_in_inference_graph(inference_nncf_graph, self.mat_mul_metatypes)
+    def get_weight_nodes(self, nncf_grpah: NNCFGraph, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
+        return get_weight_nodes(inference_nncf_graph, self.mat_mul_metatypes)
 
     def is_matmul_with_constant(self, node: NNCFNode, nncf_graph: NNCFGraph) -> bool:
-        return node.metatype in self.mat_mul_metatypes and len(get_weight_tensor_port_ids(node, nncf_graph)) > 0
+        return is_matmul_with_constant(node, nncf_graph)

nncf/quantization/algorithms/min_max/torch_fx_backend.py

@@ -38,6 +38,7 @@
 from nncf.torch.graph.graph import PTNNCFGraph
 from nncf.torch.graph.graph import PTTargetPoint
 from nncf.torch.graph.operator_metatypes import ELEMENTWISE_OPERATIONS
+from nncf.torch.graph.operator_metatypes import MATMUL_METATYPES
 from nncf.torch.graph.transformations.commands import PTSharedFnInsertionCommand
 from nncf.torch.hardware.config import PTHWConfig
 from nncf.torch.model_graph_manager import get_weight_tensor_port_ids
@@ -49,7 +50,8 @@
 from nncf.torch.quantization.layers import PTQuantizerSpec
 from nncf.torch.quantization.layers import get_scale_shape
 from nncf.torch.quantization.strip import convert_to_torch_fakequantizer
-from nncf.torch.utils import get_weight_nodes_in_inference_graph
+from nncf.torch.utils import get_weight_nodes
+from nncf.torch.utils import is_matmul_with_constant
 
 
 class FXMinMaxAlgoBackend(MinMaxAlgoBackend):
@@ -59,7 +61,7 @@ def preserved_metatypes(self) -> List[OperatorMetatype]:
 
     @property
     def mat_mul_metatypes(self) -> List[OperatorMetatype]:
-        return [om.PTLinearMetatype, om.PTMatMulMetatype]
+        return MATMUL_METATYPES
 
     @property
     def post_processing_metatypes(self) -> List[OperatorMetatype]:
@@ -305,8 +307,8 @@ def get_ignored_metatypes(model_type: ModelType, device: TargetDevice) -> List[O
     def get_ignored_names_by_layer_attributes(nncf_graph: NNCFGraph) -> Set[str]:
         return set()
 
-    def get_weight_nodes(self, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
-        return get_weight_nodes_in_inference_graph(inference_nncf_graph, self.mat_mul_metatypes)
+    def get_weight_nodes(self, nncf_grpah: NNCFGraph, inference_nncf_graph: NNCFGraph) -> List[NNCFNode]:
+        return get_weight_nodes(inference_nncf_graph, self.mat_mul_metatypes)
 
     def is_matmul_with_constant(self, node: NNCFNode, nncf_graph: NNCFGraph) -> bool:
-        return node.metatype in self.mat_mul_metatypes and len(get_weight_tensor_port_ids(node, nncf_graph)) > 0
+        return is_matmul_with_constant(node, nncf_graph)

nncf/torch/graph/operator_metatypes.py

@@ -1219,3 +1219,5 @@ def get_operator_metatypes() -> List[Type[OperatorMetatype]]:
     PTModuleEmbeddingBagMetatype,
     PTModuleEmbeddingMetatype,
 ]
+
+MATMUL_METATYPES = [PTLinearMetatype, PTMatMulMetatype, PTAddmmMetatype]

nncf/torch/utils.py

@@ -11,7 +11,7 @@
 import random
 from collections import OrderedDict
 from contextlib import contextmanager
-from typing import Any, Dict, Generator, List, Type
+from typing import Any, Dict, Generator, List
 
 import numpy as np
 import torch
@@ -31,7 +31,9 @@
 from nncf.torch.dynamic_graph.scope import Scope
 from nncf.torch.dynamic_graph.scope import ScopeElement
 from nncf.torch.dynamic_graph.trace_tensor import TracedTensorMixin
+from nncf.torch.graph.operator_metatypes import MATMUL_METATYPES
 from nncf.torch.layer_utils import _NNCFModuleMixin
+from nncf.torch.model_graph_manager import get_weight_tensor_port_ids
 from nncf.torch.structures import ExecutionParameters
 
 
@@ -472,14 +474,18 @@ def get_model_dtype(model: torch.nn.Module) -> torch.dtype:
     return dtype
 
 
-def get_weight_nodes_in_inference_graph(
-    inference_nncf_graph: NNCFGraph, mat_mul_metatypes: List[Type[om.PTOperatorMetatype]]
+def get_weight_nodes(
+    nncf_graph: NNCFGraph,
+    inference_nncf_graph: NNCFGraph,
 ) -> List[NNCFNode]:
     """
     Returns nodes that have weights.
 
     :param nncf_graph: Instance of inference NNCFGraph,
+        which contains shape of and constant subgraphs.
+    :param inference_nncf_graph: Instance of inference NNCFGraph,
         which does not contain shape of and constant subgraphs.
+
     :return: All nodes with weights.
     """
     weight_nodes_candidates = [
@@ -489,28 +495,18 @@ def get_weight_nodes_in_inference_graph(
     ]
     weight_nodes = []
     for node in weight_nodes_candidates:
-        if node.metatype in mat_mul_metatypes and not is_matmul_with_constant_in_inference_graph(
-            node, inference_nncf_graph
-        ):
-            continue
-        weight_nodes.append(node)
+        if is_matmul_with_constant(node, nncf_graph):
+            weight_nodes.append(node)
     return weight_nodes
 
 
-def is_matmul_with_constant_in_inference_graph(node: NNCFNode, inference_nncf_graph: NNCFGraph) -> bool:
+def is_matmul_with_constant(node: NNCFNode, nncf_graph: NNCFGraph) -> bool:
     """
     Determines whether the given node in the NNCF graph represents a matmul with a constant input.
 
     :param node: A NNCFNode instance.
-    :param inference_nncf_graph: An inference NNCFGraph instance.
+    :param nncf_graph: Instance of inference NNCFGraph,
+        which contains shape of and constant subgraphs.
     :return: True if given node is a matmul with a constant input, False otherwise.
     """
-    if node.metatype == om.PTLinearMetatype:
-        return True
-
-    # Inference graph does not contain constants, so
-    # any missed input edge means it is a constant branch.
-    is_matmul_metatype = node.metatype in [om.PTMatMulMetatype, om.PTAddmmMetatype]
-    inputs_missed = 1 <= len(inference_nncf_graph.get_input_edges(node)) < len(node.metatype.weight_port_ids)
-
-    return is_matmul_metatype and inputs_missed
+    return node.metatype in MATMUL_METATYPES and len(get_weight_tensor_port_ids(node, nncf_graph)) > 0