[TorchFX][FBC] Constant linear layers support (#2866)

### Changes

Constant linear layers support

### Reason for changes

To support swint_v2_s FBC

### Related tickets

#2766 

### Tests
Build post_training_quantization/444/ is finished successfully
Unit test `test_model_transformer.test_model_extraction` is presented

nncf/experimental/torch/fx/commands.py

@@ -33,5 +33,5 @@ def __init__(
         :param priority: Transformation priority.
         """
         super().__init__(TransformationType.INSERT)
-        self.tranformation_fn = transformation_fn
+        self.transformation_fn = transformation_fn
         self.priority = priority

nncf/experimental/torch/fx/model_transformer.py

@@ -18,6 +18,7 @@
 
 from nncf.common.graph.model_transformer import ModelTransformer
 from nncf.experimental.torch.fx.commands import FXApplyTransformationCommand
+from nncf.experimental.torch.fx.node_utils import get_graph_node_by_name
 from nncf.torch.graph.transformations.commands import PTModelExtractionCommand
 from nncf.torch.graph.transformations.layout import PTTransformationLayout
 
@@ -97,7 +98,23 @@ def _apply_model_extraction(
         # TODO(dlyakhov): reduce memory consumption by
         # more optimal splitting implementation.
         splitted_gm = split_by_tags(model, tags)
-        return splitted_gm.extracted
+
+        extracted_model = splitted_gm.extracted
+        graph: torch.fx.Graph = extracted_model.graph
+        # Check extracted model has inputs.
+        # It is possible to have two constant inputs
+        # for the target layer, an placeholder is being
+        # placed to the input port.
+        target_node = get_graph_node_by_name(graph, node_name)
+        input_node = target_node.all_input_nodes[0]
+        if input_node.op != "placeholder":
+            with graph.inserting_before(target_node):
+                new_input_node = graph.create_node(
+                    "placeholder", "placeholder_node", (), {}, name="placeholder_graph_node"
+                )
+            target_node.replace_input_with(input_node, new_input_node)
+        extracted_model.graph.eliminate_dead_code()
+        return extracted_model
 
     @staticmethod
     def _apply_transformation(
@@ -112,5 +129,5 @@ def _apply_transformation(
         :return: Target model after all transformations were applied.
         """
         for transformation in transformations:
-            transformation.tranformation_fn(model)
+            transformation.transformation_fn(model)
         return model

nncf/experimental/torch/fx/transformations.py

@@ -79,7 +79,7 @@ def bias_update_transformation(model: torch.fx.GraphModule):
     return bias_update_transformation
 
 
-def qdq_insertion_tranformation_builder(
+def qdq_insertion_transformation_builder(
     quantizer: FakeQuantize, target_points: List[PTTargetPoint]
 ) -> TransformationFNType:
     """
@@ -92,7 +92,7 @@ def qdq_insertion_tranformation_builder(
         inherited from the given quantizer to each given target point.
     """
 
-    def qdq_insertion_tranformation(model: torch.fx.GraphModule):
+    def qdq_insertion_transformation(model: torch.fx.GraphModule):
         if any(tp.target_type != TargetType.OPERATION_WITH_WEIGHTS for tp in target_points) and len(target_points) > 1:
             raise RuntimeError(
                 "Insertion of shared qdq pair for the weights is not supported."
@@ -101,7 +101,7 @@ def qdq_insertion_tranformation(model: torch.fx.GraphModule):
         for target_point in target_points:
             insert_one_qdq(model, target_point, quantizer)
 
-    return qdq_insertion_tranformation
+    return qdq_insertion_transformation
 
 
 def insert_one_qdq(model: torch.fx.GraphModule, target_point: PTTargetPoint, quantizer: FakeQuantize):
@@ -310,7 +310,7 @@ def _is_bn_node(node: torch.fx.Node):
 
 def fuse_conv_bn(model: torch.fx.GraphModule) -> None:
     """
-    BatchNorm operations have 3 output ports, to make it easier for alorithms to work with
+    BatchNorm operations have 3 output ports, to make it easier for algorithms to work with
     the target graph BatchNorm operations are being fused
 
     :param model: Model to apply transformations to.
@@ -342,7 +342,7 @@ def apply_quantization_transformations(model: torch.fx.GraphModule) -> None:
     :param model: Model to apply transformations to.
     """
     # BatchNorm operations have 3 output ports,
-    # to make it easier for alorithms to work
+    # to make it easier for algorithms to work
     # with the target graph BatchNorm operations
     # are being fused
     fuse_conv_bn(model)
@@ -484,7 +484,7 @@ def _merge_node_and_bias(model: torch.fx.GraphModule, is_target_node: Callable[[
     Check which node should be merged by the given `is_target_node` predicate.
 
     :param model: Target model.
-    :param is_target_node: Predicate to specify nodes which shoudld be merged with the bias
+    :param is_target_node: Predicate to specify nodes which should be merged with the bias
     """
     add_node_targets = (torch.ops.aten.add_.Tensor,)
     for n in model.graph.nodes:

nncf/quantization/algorithms/min_max/torch_fx_backend.py

@@ -28,7 +28,7 @@
 from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
 from nncf.experimental.common.tensor_statistics.statistics import MinMaxTensorStatistic
 from nncf.experimental.torch.fx.commands import FXApplyTransformationCommand
-from nncf.experimental.torch.fx.transformations import qdq_insertion_tranformation_builder
+from nncf.experimental.torch.fx.transformations import qdq_insertion_transformation_builder
 from nncf.parameters import ModelType
 from nncf.parameters import TargetDevice
 from nncf.quantization.advanced_parameters import StatisticsType
@@ -288,7 +288,7 @@ def create_quantizer_insertion_command(
         quantizer = FXMinMaxAlgoBackend._create_quantizer(
             quantizer_config, scale_shape, parameters, target_point.target_type
         )
-        transformation = qdq_insertion_tranformation_builder(quantizer, [target_point])
+        transformation = qdq_insertion_transformation_builder(quantizer, [target_point])
         return FXApplyTransformationCommand(transformation)
 
     @staticmethod
@@ -308,7 +308,7 @@ def create_unified_scales_quantizers_insertion_commands(
 
         transformations = []
         for tp in target_points:
-            transformation = qdq_insertion_tranformation_builder(quantizer, [tp])
+            transformation = qdq_insertion_transformation_builder(quantizer, [tp])
             transformations.append(FXApplyTransformationCommand(transformation))
         return transformations
 

nncf/torch/graph/transformations/serialization.py

@@ -29,7 +29,7 @@ def serialize_transformations(transformations_layout: TransformationLayout) -> D
     """
     Serializes given transformation layout to a dict.
 
-    :param tranformation_layout: Given transformation layout.
+    :param transformation_layout: Given transformation layout.
     :return: Serialized representation of given transformation layout as a dict.
     """
     transformation_commands = []

tests/post_training/data/ptq_reference_data.yaml

@@ -34,6 +34,12 @@ torchvision/resnet18_backend_CUDA_TORCH:
   metric_value: 0.69152
 torchvision/resnet18_backend_FX_TORCH:
   metric_value: 0.6946
+torchvision/swin_v2_s_backend_FP32:
+  metric_value: 0.83712
+torchvision/swin_v2_s_backend_OV:
+  metric_value: 0.83638
+torchvision/swin_v2_s_backend_FX_TORCH:
+  metric_value: 0.82908
 timm/crossvit_9_240_backend_CUDA_TORCH:
   metric_value: 0.689
 timm/crossvit_9_240_backend_FP32:

tests/post_training/model_scope.py

@@ -75,6 +75,17 @@
         "backends": [BackendType.FX_TORCH, BackendType.TORCH, BackendType.CUDA_TORCH, BackendType.OV, BackendType.ONNX],
         "batch_size": 128,
     },
+    {
+        "reported_name": "torchvision/swin_v2_s",
+        "model_id": "swin_v2_s",
+        "pipeline_cls": ImageClassificationTorchvision,
+        "compression_params": {
+            "model_type": ModelType.TRANSFORMER,
+            "advanced_parameters": AdvancedQuantizationParameters(smooth_quant_alpha=0.5),
+        },
+        "backends": [BackendType.FX_TORCH, BackendType.OV],
+        "batch_size": 1,
+    },
     # Timm models
     {
         "reported_name": "timm/crossvit_9_240",

tests/post_training/pipelines/image_classification_torchvision.py

@@ -43,13 +43,14 @@ def prepare_model(self) -> None:
         model = model_cls(weights=self.model_weights)
         model.eval()
 
-        self.static_input_size = [self.batch_size, 3, 224, 224]
+        default_input_size = [self.batch_size, 3, 224, 224]
+        self.dummy_tensor = self.model_weights.transforms()(torch.rand(default_input_size))
+        self.static_input_size = list(self.dummy_tensor.shape)
+
         self.input_size = self.static_input_size.copy()
         if self.batch_size > 1:  # Dynamic batch_size shape export
             self.input_size[0] = -1
 
-        self.dummy_tensor = torch.rand(self.static_input_size)
-
         if self.backend == BackendType.FX_TORCH:
             with torch.no_grad():
                 with disable_patching():

tests/torch/data/reference_graphs/fx/extracted/ConstantModelExtractionModel.dot

@@ -0,0 +1,11 @@
+strict digraph  {
+"0 _conv_w" [id=0, type=get_attr];
+"1 add" [id=1, type=add];
+"2 placeholder_graph_node" [id=2, type=input];
+"3 conv2d" [id=3, type=conv2d];
+"4 output" [id=4, type=output];
+"0 _conv_w" -> "1 add";
+"1 add" -> "3 conv2d";
+"2 placeholder_graph_node" -> "3 conv2d";
+"3 conv2d" -> "4 output";
+}

tests/torch/data/reference_graphs/fx/extracted/ModelExtractionModel.dot

@@ -0,0 +1,11 @@
+strict digraph  {
+"0 _conv_w" [id=0, type=get_attr];
+"1 add" [id=1, type=add];
+"2 arg0_1" [id=2, type=input];
+"3 conv2d" [id=3, type=conv2d];
+"4 output" [id=4, type=output];
+"0 _conv_w" -> "1 add";
+"1 add" -> "3 conv2d";
+"2 arg0_1" -> "3 conv2d";
+"3 conv2d" -> "4 output";
+}

tests/torch/fx/test_model_transformer.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2024 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Tuple
+
+import pytest
+import torch
+from torch._export import capture_pre_autograd_graph
+
+from nncf.common.graph.transformations.layout import TransformationLayout
+from nncf.experimental.torch.fx.model_transformer import FXModelTransformer
+from nncf.experimental.torch.fx.nncf_graph_builder import GraphConverter
+from nncf.torch import disable_patching
+from nncf.torch.graph.transformations.commands import PTModelExtractionCommand
+from tests.torch.test_compressed_graph import check_graph
+from tests.torch.test_models.synthetic import ConvolutionWithAllConstantInputsModel
+from tests.torch.test_models.synthetic import ConvolutionWithNotTensorBiasModel
+
+
+@dataclass
+class ModelExtractionTestCase:
+    model: torch.nn.Module
+    input_shape: Tuple[int, ...]
+    command: PTModelExtractionCommand
+    ref: None = None
+
+
+EXTRACTED_GRAPHS_DIR_NAME = Path("fx") / "extracted"
+
+MODEL_EXTRACTION_CASES = (
+    ModelExtractionTestCase(
+        ConvolutionWithNotTensorBiasModel, (1, 1, 3, 3), PTModelExtractionCommand(["conv2d"], ["conv2d"])
+    ),
+    ModelExtractionTestCase(
+        ConvolutionWithAllConstantInputsModel, (1, 1, 3, 3), PTModelExtractionCommand(["conv2d"], ["conv2d"])
+    ),
+)
+
+
+def idfn(value: Any):
+    if isinstance(value, ModelExtractionTestCase):
+        return value.model.__name__
+    return None
+
+
+@pytest.mark.parametrize("test_case", MODEL_EXTRACTION_CASES, ids=idfn)
+def test_model_extraction(test_case: ModelExtractionTestCase):
+    with torch.no_grad():
+        with disable_patching():
+            captured_model = capture_pre_autograd_graph(test_case.model(), (torch.ones(test_case.input_shape),))
+    layout = TransformationLayout()
+    layout.register(test_case.command)
+    extracted_model = FXModelTransformer(captured_model).transform(layout)
+    nncf_graph = GraphConverter.create_nncf_graph(extracted_model)
+    check_graph(nncf_graph, f"{test_case.model.__name__}.dot", str(EXTRACTED_GRAPHS_DIR_NAME))

tests/torch/fx/test_models.py

@@ -9,7 +9,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import json
 import os
 from dataclasses import dataclass
@@ -39,8 +38,8 @@
 from tests.torch import test_models
 from tests.torch.test_compressed_graph import check_graph
 
-FX_DIR_NAME = "fx"
-FX_QUANTIZED_DIR_NAME = "fx/quantized"
+FX_DIR_NAME = Path("fx")
+FX_QUANTIZED_DIR_NAME = Path("fx") / "quantized"
 
 
 @dataclass

tests/torch/test_models/synthetic.py

@@ -501,3 +501,24 @@ def forward(self, x):
         unbinded_processed[0] = self.conv4(y_unbinded[0])
         y = torch.cat(unbinded_processed, axis=0)
         return y
+
+
+class ConvolutionWithNotTensorBiasModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._conv_w = nn.Parameter(torch.ones((1, 1, 1, 1)))
+
+    def forward(self, x):
+        w = self._conv_w + 10
+        return nn.functional.conv2d(x, w)
+
+
+class ConvolutionWithAllConstantInputsModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._conv_w = nn.Parameter(torch.ones((1, 1, 1, 1)))
+        self._conv_i = nn.Parameter(torch.ones((1, 1, 1, 1)))
+
+    def forward(self, x):
+        w = self._conv_w + 10
+        return x + nn.functional.conv2d(self._conv_i, w)