[PT] FBC for linear (#3808)

### Changes

Add support linear operation for FBC 

### Related tickets

CVS-111111

### Tests

https://github.com/openvinotoolkit/nncf/actions/runs/20462273440
manual/job/post_training_quantization/766/

Author: AlexanderDokuchaev

src/nncf/torch/function_hook/extractor.py

@@ -127,7 +127,7 @@ def extract_conv(
     """
     Extracts a convolutional layer from an NNCF graph and constructs an ExtractedFunc module.
 
-    :param model: The NNCF network containing the layer.
+    :param model: The nn.Module containing the layer.
     :param graph: The NNCF graph.
     :param input_nodes: The name of input node.
     :param output_nodes: The name of output node.
@@ -145,13 +145,13 @@ def extract_conv(
 
     weight_node = get_const_node(input_node, 1, graph)
     if weight_node is None:
-        msg = "Weight node not found for {input_node}"
+        msg = f"Weight node not found for {input_node}"
         raise nncf.InternalError(msg)
     weight = get_const_data(weight_node, model)
 
     hook_storage = get_hook_storage(model)
     with torch.no_grad():
-        # Calculate weight after execution all hook fro weight data
+        # Calculate weight after execution all hook for weight data
         weight = hook_storage.execute_post_function_hooks(weight_node.node_name, 0, weight)
         weight = hook_storage.execute_pre_function_hooks(input_node.node_name, 1, weight)
 
@@ -189,17 +189,66 @@ def extract_conv(
     return nn.Sequential(conv_module, bn_module)
 
 
+def extract_linear(
+    model: nn.Module,
+    graph: PTNNCFGraph,
+    input_node: NNCFNode,
+    output_node: NNCFNode,
+) -> ExtractedFunc:
+    """
+    Extracts a linear layer from an NNCF graph and constructs an ExtractedFunc module.
+
+    :param model: The nn.Module containing the layer.
+    :param graph: The NNCF graph.
+    :param input_node: The name of input node.
+    :param output_node: The name of output node.
+    :return: The extracted linear layer as an ExtractedFunc module.
+    """
+    if input_node != output_node:
+        msg = "Only one input and output node supported."
+        raise nncf.InternalError(msg)
+
+    layer_attrs = input_node.layer_attributes
+
+    if not isinstance(layer_attrs, PT2OpLayerAttributes):
+        msg = f"Expected PT2OpLayerAttributes for input_node.layer_attributes, actual: {type(layer_attrs)}"
+        raise nncf.InternalError(msg)
+
+    weight_node = get_const_node(input_node, 1, graph)
+    if weight_node is None:
+        msg = f"Weight node not found for {input_node}"
+        raise nncf.InternalError(msg)
+    weight = get_const_data(weight_node, model)
+
+    hook_storage = get_hook_storage(model)
+    with torch.no_grad():
+        # Calculate weight after execution all hook for weight data
+        weight = hook_storage.execute_post_function_hooks(weight_node.node_name, 0, weight)
+        weight = hook_storage.execute_pre_function_hooks(input_node.node_name, 1, weight)
+
+    bias_node = get_const_node(input_node, 2, graph)
+    bias = get_const_data(bias_node, model) if bias_node is not None else None
+
+    layer_kwarg = {
+        "weight": weight,
+        "bias": bias,
+    }
+    linear_module = ExtractedFunc(layer_attrs.func, layer_kwarg)
+    return linear_module
+
+
 def extract_model(
     model: nn.Module, graph: PTNNCFGraph, input_nodes: list[str], output_nodes: list[str]
 ) -> Optional[nn.Module]:
     """
-    Extracts a submodule from a given NNCF network containing only the nodes from the input to the output node.
+    Extracts a submodule from a given nn.Module containing only the nodes from the input to the output node.
 
     Supported subgraph:
       - Conv
       - Conv + BatchNorm
+      - Linear
 
-    :param model: The NNCF network to extract the submodule from.
+    :param model: The nn.Module to extract the submodule from.
     :param input_nodes: List containing names of the input nodes for the submodule.
     :param output_nodes: List containing names of the output nodes for the submodule.
     :return: An nn.Module containing the extracted submodel, or None if extraction is not supported.
@@ -214,5 +263,8 @@ def extract_model(
     if input_node.metatype in CONV_METATYPES:
         return extract_conv(model, graph, input_node, output_node)
 
-    nncf_logger.debug(f"Can`t extract module for {input_node.node_name}")
+    if input_node.metatype is om.PTLinearMetatype:
+        return extract_linear(model, graph, input_node, output_node)
+
+    nncf_logger.debug(f"Can not extract module for {input_node.node_name}")
     return None

src/nncf/torch/graph/operator_metatypes.py

@@ -287,6 +287,7 @@ class PTLinearMetatype(PTOperatorMetatype):
     output_channel_axis = -1
     num_expected_input_edges = 2
     weight_port_ids = [1]
+    bias_port_id = 2
 
 
 @PT_OPERATOR_METATYPES.register()

src/nncf/torch/model_graph_manager.py

@@ -35,7 +35,7 @@
     om.PTConvTranspose3dMetatype,
 ]
 
-OPERATORS_WITH_BIAS_METATYPES = CONV_META_TYPES
+OPERATORS_WITH_BIAS_METATYPES = CONV_META_TYPES + [om.PTLinearMetatype]
 CONV_FUSED_META_TYPES = [om.PTBatchNormMetatype]
 
 

tests/cross_fw/test_templates/helpers.py

@@ -163,13 +163,13 @@ def forward(self, x):
 
 
 class FCTestModel(nn.Module):
-    INPUT_SIZE = [1, 1, 4, 4]
+    INPUT_SIZE = [1, 1, 3, 3]
 
     def __init__(self):
         super().__init__()
-        self.fc = nn.Linear(4, 2)
-        self.fc.weight.data = torch.Tensor([[0.1, 0.2, 0.3, 0.2], [0.3, -0.1, 0.2, 0.4]])
-        self.fc.bias.data = torch.Tensor([1.0, 1.1])
+        self.fc = nn.Linear(3, 2)
+        self.fc.weight.data = torch.Tensor([[0.1, 0.2, 0.3], [0.3, -0.1, 0.2]])
+        self.fc.bias.data = torch.Tensor([1.0, 2.0])
 
     def forward(self, x):
         x = self.fc(x)

tests/cross_fw/test_templates/test_fast_bias_correction.py

@@ -10,6 +10,8 @@
 # limitations under the License.
 
 from abc import abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
 from typing import TypeVar
 
 import pytest
@@ -22,12 +24,22 @@
 from nncf.quantization.algorithms.post_training.algorithm import PostTrainingQuantization
 from tests.cross_fw.test_templates.helpers import ConvBNTestModel
 from tests.cross_fw.test_templates.helpers import ConvTestModel
+from tests.cross_fw.test_templates.helpers import FCTestModel
 from tests.cross_fw.test_templates.helpers import get_static_dataset
 
 TModel = TypeVar("TModel")
 TTensor = TypeVar("TTensor")
 
 
+@dataclass
+class TestCase:
+    model_cls: type
+    ref_bias: list
+
+    def __str__(self):
+        return self.model_cls.__name__
+
+
 class TemplateTestFBCAlgorithm:
     @staticmethod
     @abstractmethod
@@ -104,18 +116,19 @@ def get_quantization_algorithm():
         )
 
     @pytest.mark.parametrize(
-        "model_cls, ref_bias",
+        "params",
         (
-            (ConvTestModel, [0.0288348, 1.0838453]),
-            (ConvBNTestModel, [0.08396978, 1.1676897]),
+            TestCase(ConvTestModel, [0.0288348, 1.0838453]),
+            TestCase(ConvBNTestModel, [0.08396978, 1.1676897]),
+            TestCase(FCTestModel, [0.9999, 1.9989]),
         ),
+        ids=str,
     )
-    def test_update_bias(self, model_cls, ref_bias, tmpdir):
-        model = self.backend_specific_model(model_cls(), tmpdir)
-        dataset = get_static_dataset(model_cls.INPUT_SIZE, self.get_transform_fn(), self.fn_to_type)
+    def test_update_bias(self, params: TestCase, tmpdir: Path):
+        model = self.backend_specific_model(params.model_cls(), tmpdir)
+        dataset = get_static_dataset(params.model_cls.INPUT_SIZE, self.get_transform_fn(), self.fn_to_type)
 
         quantization_algorithm = self.get_quantization_algorithm()
         graph = NNCFGraphFactory.create(model)
         quantized_model = quantization_algorithm.apply(model, graph, dataset=dataset)
-
-        self.check_bias(quantized_model, ref_bias)
+        self.check_bias(quantized_model, params.ref_bias)

tests/torch/test_model_graph_manager.py

@@ -111,7 +111,7 @@ def test_get_potential_fused_node(self, model_desc):
         "ConvBiasBNTestModel": True,
         "ConvBNTestModel": True,
         "ConvTestModel": True,
-        "FCTestModel": False,
+        "FCTestModel": True,
         "MultipleConvTestModel": True,
         "CustomConvTestModel": True,
         "CustomConvBNTestModel": True,
@@ -152,7 +152,7 @@ def test_get_const_node(self, model_desc, port_id):
             [[[[0.1000, -2.0000], [1.0000, 0.1000]]], [[[0.1000, 2.0000], [-1.0000, 0.1000]]]],
             [0.1000, 1.0000],
         ),
-        "FCTestModel": ([[0.1000, 0.2000, 0.3000, 0.2000], [0.3000, -0.1000, 0.2000, 0.4000]], [1.0000, 1.1000]),
+        "FCTestModel": ([[0.1000, 0.2000, 0.3000], [0.3000, -0.1000, 0.2000]], [1.0000, 2.0000]),
         "MultipleConvTestModel": (
             [[[[-2.4661, 0.3623], [0.3765, -0.1808]]], [[[0.3930, 0.4327], [-1.3627, 1.3564]]]],
             [0.6688, -0.7077],

tests/torch2/function_hook/test_extractor.py

@@ -48,6 +48,11 @@
         "conv/conv2d/0",
         "conv/conv2d/0",
     ),
+    (
+        helpers.FCTestModel,
+        "fc/linear/0",
+        "fc/linear/0",
+    ),
 )