[TorchFX][FBC] Constant linear layers support

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 a linear 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(

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,57 @@
+# 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 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 ConstantModelExtractionModel
+from tests.torch.test_models.synthetic import ModelExtractionModel
+
+
+@dataclass
+class ModelExtractionTestCase:
+    model: torch.nn.Module
+    input_shape: Tuple[int, ...]
+    command: PTModelExtractionCommand
+    ref: None = None
+
+
+EXTRAXCTED_GRAPHS_DIR_NAME = Path("fx") / "extracted"
+
+MODEL_EXTRACTION_CASES = (
+    ModelExtractionTestCase(ModelExtractionModel, (1, 1, 3, 3), PTModelExtractionCommand(["conv2d"], ["conv2d"])),
+    ModelExtractionTestCase(
+        ConstantModelExtractionModel, (1, 1, 3, 3), PTModelExtractionCommand(["conv2d"], ["conv2d"])
+    ),
+)
+
+
+@pytest.mark.parametrize("test_case", MODEL_EXTRACTION_CASES)
+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", EXTRAXCTED_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 ModelExtractionModel(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 ConstantModelExtractionModel(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)