[TorchFX] Deptwise convolution support (#2896)

### Changes

Depthwise convolution subtypes are integrated to torchFX nncf graph
builder

### Reason for changes

To correctly quantize deptiwise convolutions

### Related tickets

#2766 

### Tests

`tests/torch/fx/test_models.py` is updated

nncf/experimental/torch/fx/nncf_graph_builder.py

@@ -15,9 +15,13 @@
 
 import nncf.torch.graph.operator_metatypes as om
 from nncf.common.graph import NNCFNode
+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 UnknownMetatype
 from nncf.common.logging import nncf_logger
+from nncf.experimental.torch.fx.node_utils import get_tensor_constant_from_node
+from nncf.torch.dynamic_graph.layer_attributes_handlers import apply_args_defaults
 from nncf.torch.graph.graph import PTNNCFGraph
 from nncf.torch.graph.operator_metatypes import PT_OPERATOR_METATYPES
 
@@ -27,12 +31,48 @@ class GraphConverter:
     Builds the NNCFGraph from an torch.fx.GraphModule instance.
     """
 
+    def _get_layer_attributes(
+        node: torch.fx.Node, metatype: om.OperatorMetatype, model: torch.fx.GraphModule
+    ) -> BaseLayerAttributes:
+        """
+        Collects layer attributes for the given node.
+
+        :param node: Given node.
+        :param metatype: Given node metatype.
+        :param model: Target GraphModule instance.
+        :return: Given node layer attributes.
+        """
+        if metatype in [om.PTConv1dMetatype, om.PTConv2dMetatype, om.PTConv3dMetatype]:
+            conv_default_args = [(arg.name, arg.default_value) for arg in node.target._schema.arguments]
+            kwargs = apply_args_defaults(node.args, node.kwargs, conv_default_args)
+
+            weight_node = kwargs["weight"]
+            if weight_node.op != "get_attr":
+                # Convs with constant subgraphs or two inputs are not supported yet.
+                return None
+            weight = get_tensor_constant_from_node(weight_node, model)
+            return ConvolutionLayerAttributes(
+                weight_requires_grad=False,
+                in_channels=weight.shape[0],
+                out_channels=weight.shape[1],
+                kernel_size=list(weight.shape[2:]),
+                stride=kwargs["stride"],
+                dilations=kwargs["dilation"],
+                groups=kwargs["groups"],
+                padding_values=kwargs["padding"],
+                transpose=False,
+            )
+        return None
+
     @staticmethod
-    def _get_node_type_and_metatype(node: torch.fx.Node) -> Tuple[str, om.OperatorMetatype]:
+    def _get_node_type_and_metatype(
+        node: torch.fx.Node, model: torch.fx.GraphModule
+    ) -> Tuple[str, om.OperatorMetatype]:
         """
         Retrieves node's type and metatype.
 
         :param node: Given node.
+        :param model: Given GraphModule.
         :return: Node's type and metatype.
         """
         if node.op == "placeholder":
@@ -58,6 +98,11 @@ def _get_node_type_and_metatype(node: torch.fx.Node) -> Tuple[str, om.OperatorMe
             node_metatype = UnknownMetatype
         if node_metatype is UnknownMetatype:
             nncf_logger.debug(f"Unknown metatype for node: {node}")
+
+        if node_metatype.get_subtypes():
+            layer_attrs = GraphConverter._get_layer_attributes(node, node_metatype, model)
+            node_subtype = node_metatype.determine_subtype(layer_attrs)
+            node_metatype = node_subtype or node_metatype
         return node_type, node_metatype
 
     @staticmethod
@@ -74,7 +119,7 @@ def create_nncf_graph(model: torch.fx.GraphModule) -> PTNNCFGraph:
         nncf_graph = PTNNCFGraph()
 
         for source_node in model.graph.nodes:
-            node_type, node_metatype = GraphConverter._get_node_type_and_metatype(source_node)
+            node_type, node_metatype = GraphConverter._get_node_type_and_metatype(source_node, model)
 
             nncf_graph.add_nncf_node(
                 node_name=source_node.name,

tests/torch/fx/test_models.py

@@ -90,7 +90,7 @@ def get_ref_metatypes_from_json(
         if not os.path.exists(json_parent_dir):
             os.makedirs(json_parent_dir)
         with safe_open(complete_path, "w") as file:
-            json.dump(model_metatypes, file)
+            json.dump(model_metatypes, file, indent=4)
 
     with safe_open(complete_path, "r") as file:
         return json.load(file)
@@ -115,7 +115,7 @@ def test_model(test_case: ModelCase):
         check_graph(nncf_graph, dot_filename, FX_DIR_NAME)
 
         # Check metatypes
-        model_metatypes = {n.node_name: n.metatype.name for n in nncf_graph.get_all_nodes()}
+        model_metatypes = {n.node_name: n.metatype.__name__ for n in nncf_graph.get_all_nodes()}
         ref_metatypes = get_ref_metatypes_from_json(model_name, model_metatypes)
         assert model_metatypes == ref_metatypes