mct-nightly 2.4.0.20250925.543__py3-none-any.whl → 2.4.2.20250926.532__py3-none-any.whl

model_compression_toolkit/core/pytorch/back2framework/float_model_builder.py

@@ -24,6 +24,8 @@ from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder, \
     PytorchModel
 
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+
 
 class FloatPyTorchModel(PytorchModel):
     """
@@ -32,16 +34,19 @@ class FloatPyTorchModel(PytorchModel):
 
     def __init__(self,
                  graph: common.Graph,
-                 append2output=None):
+                 append2output=None,
+                 fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO):
         """
 
         Args:
             graph: Graph to build its corresponding Pytorch model.
             append2output: List of nodes or OutTensor objects.
+            fw_info: Framework information (e.g., mapping from layers to their attributes to quantize).
         """
 
         super().__init__(graph,
-                         append2output)
+                         append2output,
+                         fw_info)
 
     def _quantize_node_activations(self,
                                    node: BaseNode,
@@ -66,17 +71,20 @@ class FloatPyTorchModelBuilder(PyTorchModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
+                 fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
+            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
         """
 
         super().__init__(graph,
                          append2output,
+                         fw_info,
                          return_float_outputs)
 
     def build_model(self) -> Tuple[PytorchModel, UserInformation]:
@@ -86,4 +94,5 @@ class FloatPyTorchModelBuilder(PyTorchModelBuilder):
 
         """
         return FloatPyTorchModel(self.graph,
-                                 self.append2output), self.graph.user_info
+                                 self.append2output,
+                                 self.fw_info), self.graph.user_info

model_compression_toolkit/core/pytorch/back2framework/mixed_precision_model_builder.py

@@ -23,6 +23,7 @@ from model_compression_toolkit.core import FrameworkInfo, common
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 from model_compression_toolkit.core.pytorch.mixed_precision.configurable_activation_quantizer import \
     ConfigurableActivationQuantizer
 from model_compression_toolkit.core.pytorch.mixed_precision.configurable_weights_quantizer import \
@@ -37,12 +38,14 @@ class MixedPrecisionPyTorchModelBuilder(PyTorchModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
+                 fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
+            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
         """
 
@@ -50,6 +53,7 @@ class MixedPrecisionPyTorchModelBuilder(PyTorchModelBuilder):
 
         super().__init__(graph,
                          append2output,
+                         fw_info,
                          return_float_outputs,
                          wrapper=self.mixed_precision_wrapper,
                          get_activation_quantizer_holder_fn=self.mixed_precision_activation_holder)
@@ -73,16 +77,17 @@ class MixedPrecisionPyTorchModelBuilder(PyTorchModelBuilder):
             ValueError: if kernel attribute is quantized but not configurable.
         """
 
-        if n.kernel_attr is None or not n.is_weights_quantization_enabled(n.kernel_attr):
+        kernel_attr = self.fw_info.get_kernel_op_attributes(n.type)[0]
+        if kernel_attr is None or not n.is_weights_quantization_enabled(kernel_attr):
             return layer
-        if not n.is_configurable_weight(n.kernel_attr):  # pragma: no cover
+        if not n.is_configurable_weight(kernel_attr):  # pragma: no cover
             raise ValueError(f'Weight wrapper is not expected to be created for non-configurable weight of node {n}.')
         return PytorchQuantizationWrapper(layer,
                                           weights_quantizers={
-                                              n.kernel_attr: ConfigurableWeightsQuantizer(
+                                              kernel_attr: ConfigurableWeightsQuantizer(
                                                   **self._get_weights_configurable_quantizer_kwargs(n,
-                                                                                                    n.kernel_attr),
-                                                  kernel_attr=n.kernel_attr)})
+                                                                                                    kernel_attr),
+                                                  kernel_attr=kernel_attr)})
 
     def _get_weights_configurable_quantizer_kwargs(self, n: BaseNode, attr: str) -> Dict[str, Any]:
         """
@@ -142,13 +147,14 @@ class MixedPrecisionPyTorchModelBuilder(PyTorchModelBuilder):
         # activation number of bits (in reversed order).
         # since only kernel attribute is quantized in weights mixed precision,
         # if the node doesn't have a kernel attribute, we only sort by activation_n_bits.
-        n.sort_node_candidates()
+        n.sort_node_candidates(self.fw_info)
 
         max_candidate_idx = n.find_max_candidate_index()
 
+        kernel_attr = self.fw_info.get_kernel_op_attributes(n.type)[0]
         activation_quantizers = [ConfigurableActivationQuantizer(**{'node_q_cfg': node_q_cfg_candidates,
                                                                     'max_candidate_idx': max_candidate_idx,
-                                                                    'kernel_attr': n.kernel_attr})] \
+                                                                    'kernel_attr': kernel_attr})] \
                                 * num_of_outputs
 
         # Holder by definition uses a single quantizer for the activation quantization
@@ -171,7 +177,7 @@ class MixedPrecisionPyTorchModelBuilder(PyTorchModelBuilder):
         # creating a mapping between graph nodes and model's layers for mixed precision configurability
         model_layers = dict(model.named_children())
         conf_node2layers = {n.name: self._find_layers_in_model_by_node(n, model_layers)
-                            for n in self.graph.get_configurable_sorted_nodes()}
+                            for n in self.graph.get_configurable_sorted_nodes(self.fw_info)}
 
         return model, user_info, conf_node2layers
 
@@ -224,7 +230,8 @@ class MixedPrecisionPyTorchModelBuilder(PyTorchModelBuilder):
 
         """
         # Only layers with kernel op are considered weights configurable
-        weights_quant = False if n.kernel_attr is None else n.is_weights_quantization_enabled(n.kernel_attr)
+        kernel_attr = self.fw_info.get_kernel_op_attributes(n.type)[0]
+        weights_quant = False if kernel_attr is None else n.is_weights_quantization_enabled(kernel_attr)
         act_quant = n.is_activation_quantization_enabled()
 
         if weights_quant and not act_quant:

model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -30,6 +30,7 @@ from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.pytorch.back2framework.instance_builder import node_builder
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlaceHolder
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from mct_quantizers.common.constants import ACTIVATION_HOLDER_QUANTIZER
@@ -363,7 +364,7 @@ class PytorchModel(torch.nn.Module):
         """
         node_to_output_tensors_dict = dict()
         node_to_output_tensors_dict_float = dict()
-        configurable_nodes = self.graph.get_configurable_sorted_nodes_names()
+        configurable_nodes = self.graph.get_configurable_sorted_nodes_names(DEFAULT_PYTORCH_INFO)
         for node in self.node_sort:
             op_func = self._get_op_func(node, configurable_nodes)
             input_tensors = _build_input_tensors_list(node,
@@ -439,6 +440,7 @@ class PyTorchModelBuilder(BaseModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
+                 fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                  return_float_outputs: bool = False,
                  wrapper: Callable = None,
                  get_activation_quantizer_holder_fn: Callable = None):
@@ -447,6 +449,7 @@ class PyTorchModelBuilder(BaseModelBuilder):
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
+            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
             wrapper: A function wrapper Pytorch Layers.
             get_activation_quantizer_holder_fn: Function to retrieve a quantization holder for a node.
@@ -454,6 +457,7 @@ class PyTorchModelBuilder(BaseModelBuilder):
 
         super().__init__(graph,
                          append2output,
+                         fw_info,
                          return_float_outputs)
 
         self.wrapper = wrapper

model_compression_toolkit/core/pytorch/back2framework/quantization_wrapper/quantized_layer_wrapper.py

@@ -21,6 +21,7 @@ from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.pytorch.back2framework.quantization_wrapper.wrapper_quantize_config import \
     WrapperQuantizeConfig
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 from model_compression_toolkit.core.pytorch.utils import set_model, to_torch_tensor
 
 
@@ -92,7 +93,7 @@ class QuantizedLayerWrapper(torch.nn.Module):
             self.layer = n.type(**framework_attr)
             self.layer.load_state_dict({k: torch.Tensor(v) for k, v in n.weights.items()}, strict=False)
 
-    def _quantize_weights(self, n: BaseNode):
+    def _quantize_weights(self, n:BaseNode):
         """
         Quantize node's weights and load them as the layer's weights.
 
@@ -103,7 +104,7 @@ class QuantizedLayerWrapper(torch.nn.Module):
             None.
         """
 
-        self.weight_attrs = [n.kernel_attr]
+        self.weight_attrs = DEFAULT_PYTORCH_INFO.get_kernel_op_attributes(n.type)
 
         # float_weights is a list of weights for each attribute that we want to quantize.
         float_weights = [n.get_weights_by_keys(attr) for attr in self.weight_attrs]

model_compression_toolkit/core/pytorch/back2framework/quantized_model_builder.py

@@ -17,13 +17,13 @@ from typing import List, Tuple
 
 import torch
 
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
-from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_activation_quantization_fn
-from model_compression_toolkit.core.pytorch.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder, \
     PytorchModel
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 
 
 class QuantizedPyTorchModel(PytorchModel):
@@ -61,9 +61,7 @@ class QuantizedPyTorchModel(PytorchModel):
         if node.is_activation_quantization_enabled():
             if isinstance(input_tensors, list):
                 input_tensors = torch.cat(input_tensors, dim=0)
-            activation_quantizer = get_activation_quantization_fn(node.final_activation_quantization_cfg,
-                                                                  get_activation_quantization_fn_factory)
-            return activation_quantizer(input_tensors)
+            return node.final_activation_quantization_cfg.quantize_node_output(input_tensors)
         return input_tensors
 
 
@@ -72,17 +70,20 @@ class QuantizedPyTorchModelBuilder(PyTorchModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
+                 fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
+            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
         """
 
         super().__init__(graph,
                          append2output,
+                         fw_info,
                          return_float_outputs)
 
     def build_model(self) -> Tuple[PytorchModel, UserInformation]:

model_compression_toolkit/core/pytorch/default_framework_info.py

@@ -12,101 +12,87 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Any
-from functools import wraps
-
-from torch.nn import Hardsigmoid, ReLU, ReLU6, Softmax, Sigmoid, GELU, SELU, SiLU
-from torch.nn.functional import hardsigmoid, relu, relu6, softmax, gelu, selu, silu
+from torch.nn import Hardsigmoid, ReLU, ReLU6, Softmax, Sigmoid, GELU, SELU
+from torch.nn.functional import hardsigmoid, relu, relu6, softmax, gelu, selu
 from torch.nn import Conv2d, ConvTranspose2d, Linear
 from torch import sigmoid
 
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo, set_fw_info, ChannelAxisMapping
+from model_compression_toolkit.defaultdict import DefaultDict
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo, DEFAULT_KERNEL_ATTRIBUTES
+from mct_quantizers import QuantizationMethod
 from model_compression_toolkit.constants import SOFTMAX_THRESHOLD
 from model_compression_toolkit.core.pytorch.constants import KERNEL
-
-
-class PyTorchInfo(FrameworkInfo):
-    """
-    Extra field defined to handle Activation layer functions:
-    """
-
-    """
-    Map each layer to it's weight attribute that should get quantized.
-    If a layer that is not listed here is queried, None is returned.
-    """
-    kernel_ops_attribute_mapping = {Conv2d: KERNEL,
-                                    ConvTranspose2d: KERNEL,
-                                    Linear: KERNEL}
-
-    """
-    Map a layer to its kernel's output and input channels indices.
-    Map's values are tuples of (output_channel_index, input_channel_index).
-    Default value is returned for layers that are not included.
-    """
-    kernel_channels_mapping = {Conv2d: ChannelAxisMapping(0, 1),
-                               Linear: ChannelAxisMapping(0, 1),
-                               ConvTranspose2d: ChannelAxisMapping(1, 0)}
-
-    """
-    Map a layer to its output channel axis.
-    Where axis=-1 is the last axis
-    """
-    out_channel_axis_mapping = {Conv2d: 1,
-                                Linear: -1,
-                                ConvTranspose2d: 1}
-
-    """
-    Map from an Pytorch module to its min/max output values (if known).
-    The values are used for tensor min/max values initialization.
-    """
-    _layer_min_max_mapping = {Softmax: (0, SOFTMAX_THRESHOLD),
-                              softmax: (0, SOFTMAX_THRESHOLD),
-                              Sigmoid: (0, 1),
-                              sigmoid: (0, 1),
-                              Hardsigmoid: (0, 1),
-                              hardsigmoid: (0, 1),
-                              ReLU: (0, None),
-                              relu: (0, None),
-                              ReLU6: (0, None),
-                              relu6: (0, None),
-                              GELU: (-0.17, None),
-                              gelu: (-0.17, None),
-                              SELU: (-1.76, None),
-                              selu: (-1.76, None),
-                              silu: (-0.279, None),
-                              SiLU: (-0.279, None),
-                              }
-
-    @classmethod
-    def get_kernel_channels(cls, node_type: Any) -> ChannelAxisMapping:
-        """
-        Returns node's channels mapping from kernel_channels_mapping or framework specific default value.
-        Args:
-            node_type: A node type.
-
-        Returns:
-            Node's channels mapping.
-
-        """
-        return cls.kernel_channels_mapping.get(node_type, cls._default_channel_mapping)
-
-    @classmethod
-    def get_out_channel_axis(cls, node_type: Any):
-        """
-        Returns node's output channel mapping from out_channel_axis_mapping or framework specific default value.
-        Args:
-            node_type: A node type.
-
-        Returns:
-            Node's output channel axis.
-
-        """
-        return cls.out_channel_axis_mapping.get(node_type)
-
-
-def set_pytorch_info(func):
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        set_fw_info(PyTorchInfo)
-        return func(*args, **kwargs)
-    return wrapper
+from model_compression_toolkit.core.pytorch.quantizer.fake_quant_builder import power_of_two_quantization, \
+    symmetric_quantization, uniform_quantization
+from model_compression_toolkit.core.pytorch.quantizer.lut_fake_quant import activation_lut_kmean_quantizer
+
+"""
+Map each layer to a list of its' weights attributes that should get quantized.
+If a layer that is not listed here is queried, [None] is returned.
+"""
+KERNEL_ATTRIBUTES = DefaultDict({Conv2d: [KERNEL],
+                                 ConvTranspose2d: [KERNEL],
+                                 Linear: [KERNEL]},
+                                DEFAULT_KERNEL_ATTRIBUTES)
+
+"""
+Map a layer to its kernel's output and input channels indices.
+Map's values are tuples of (output_channel_index, input_channel_index).
+Default value is returned for layers that are not included.
+"""
+DEFAULT_CHANNEL_AXIS_DICT = DefaultDict({Conv2d: (0, 1),
+                                         Linear: (0, 1),
+                                         ConvTranspose2d: (1, 0)},
+                                        (None, None))
+
+"""
+Map a layer to its output channel axis.
+Where axis=-1 is the last axis
+"""
+DEFAULT_OUT_CHANNEL_AXIS_DICT = DefaultDict({Conv2d: 1,
+                                             Linear: -1,
+                                             ConvTranspose2d: 1},
+                                            1)
+
+
+"""
+Map from an activation function to its min/max output values (if known).
+The values are used for tensor min/max values initialization.
+"""
+ACTIVATION2MINMAX = {}  # should be an empty dict in Pytorch
+
+"""
+Map from an Pytorch module to its min/max output values (if known).
+The values are used for tensor min/max values initialization.
+"""
+LAYER2MINMAX = {Softmax: (0, SOFTMAX_THRESHOLD),
+                softmax: (0, SOFTMAX_THRESHOLD),
+                Sigmoid: (0, 1),
+                sigmoid: (0, 1),
+                Hardsigmoid: (0, 1),
+                hardsigmoid: (0, 1),
+                ReLU: (0, None),
+                relu: (0, None),
+                ReLU6: (0, None),
+                relu6: (0, None),
+                GELU: (-0.17, None),
+                gelu: (-0.17, None),
+                SELU: (-1.76, None),
+                selu: (-1.76, None),
+                }
+
+"""
+Mapping from a QuantizationMethod to an activation quantizer function.
+"""
+ACTIVATION_QUANTIZER_MAPPING = {QuantizationMethod.POWER_OF_TWO: power_of_two_quantization,
+                                QuantizationMethod.SYMMETRIC: symmetric_quantization,
+                                QuantizationMethod.UNIFORM: uniform_quantization,
+                                QuantizationMethod.LUT_POT_QUANTIZER: activation_lut_kmean_quantizer}
+
+
+DEFAULT_PYTORCH_INFO = FrameworkInfo(ACTIVATION_QUANTIZER_MAPPING,
+                                     DEFAULT_CHANNEL_AXIS_DICT,
+                                     ACTIVATION2MINMAX,
+                                     LAYER2MINMAX,
+                                     KERNEL_ATTRIBUTES,
+                                     DEFAULT_OUT_CHANNEL_AXIS_DICT)

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/const_holder_conv.py

@@ -21,18 +21,19 @@ from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.pytorch.constants import IN_CHANNELS, OUT_CHANNELS, KERNEL_SIZE, KERNEL, BIAS
-from model_compression_toolkit.core.common.framework_info import get_fw_info
+from model_compression_toolkit.core.common import FrameworkInfo
 
 
 class FunctionalConvSubstitution(common.BaseSubstitution):
     """
     Substitute functional convolutions with Layers
     """
-    def __init__(self):
+    def __init__(self, fw_info: FrameworkInfo):
         """
         Matches a functional conv node
         """
         func_node = NodeOperationMatcher(conv2d) | NodeOperationMatcher(conv_transpose2d)
+        self.fw_info = fw_info
         super().__init__(matcher_instance=func_node)
 
     def substitute(self,
@@ -55,7 +56,7 @@ class FunctionalConvSubstitution(common.BaseSubstitution):
         else:
             Logger.critical(f'Substitution filter mismatch. Layer {func_node.type}. Must be {type(Conv2d)} or {type(ConvTranspose2d)}.')  # pragma: no cover
 
-        out_channel_index, in_channel_index = get_fw_info().get_kernel_channels(new_layer)
+        out_channel_index, in_channel_index = self.fw_info.kernel_channels_mapping.get(new_layer)
 
         # Create new node of layer convolution
         if 1 not in func_node.weights:

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/relu_bound_to_power_of_2.py

@@ -95,11 +95,11 @@ class ReLUBoundToPowerOfTwo(common.BaseSubstitution):
             else:
                 return graph
         elif non_linear_node.is_match_type(hardtanh):
-            kwargs = non_linear_node.op_call_kwargs
-            if (kwargs[HARDTANH_MIN_VAL] == 0.0) and not \
-                    (np.log2(kwargs[HARDTANH_MAX_VAL]).astype(int) - np.log2(kwargs[HARDTANH_MAX_VAL]) == 0):
-                scale_factor = kwargs[HARDTANH_MAX_VAL] / self.threshold
-                non_linear_node.functional_op.__defaults__ = (0.0, self.threshold, kwargs[INPLACE])
+            if (non_linear_node.framework_attr[HARDTANH_MIN_VAL] == 0.0) and not \
+                    (np.log2(non_linear_node.framework_attr[HARDTANH_MAX_VAL]).astype(int) -
+                     np.log2(non_linear_node.framework_attr[HARDTANH_MAX_VAL]) == 0):
+                scale_factor = non_linear_node.framework_attr[HARDTANH_MAX_VAL] / self.threshold
+                non_linear_node.functional_op.__defaults__ = (0.0, self.threshold, non_linear_node.framework_attr[INPLACE])
             else:
                 return graph
         else:

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/scale_equalization.py

@@ -46,15 +46,17 @@ class ScaleEqualization(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig):
+                 quant_config: QuantizationConfig,
+                 fw_info: FrameworkInfo):
         """
         Initialize a ScaleEqualization object.
         Args:
             quant_config: Quantization configuration.
+            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, matcher_instance=MATCHER,
+        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER,
                          kernel_str=KERNEL, bias_str=BIAS)
 
 
@@ -64,13 +66,15 @@ class ScaleEqualizationWithPad(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig):
+                 quant_config: QuantizationConfig,
+                 fw_info: FrameworkInfo):
         """
         Initialize a ScaleEqualization object.
         Args:
             quant_config: Quantization configuration.
+            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, matcher_instance=MATCHER_WITH_PAD,
+        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER_WITH_PAD,
                          kernel_str=KERNEL, bias_str=BIAS)

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/shift_negative_activation.py

@@ -29,7 +29,6 @@ from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import EdgeMatcher
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.core.common.substitutions.shift_negative_activation import apply_shift_negative_correction
-from model_compression_toolkit.core.pytorch.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.core.pytorch.constants import PAD, VALUE, PADDING, BIAS, USE_BIAS
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 
@@ -215,13 +214,15 @@ def is_padding_node_and_node_has_padding(pad_node_to_consider: BaseNode,
 
 
 def pytorch_apply_shift_negative_correction(graph: Graph,
-                                            core_config: CoreConfig) -> Graph:
+                                            core_config: CoreConfig,
+                                            fw_info: FrameworkInfo) -> Graph:
     """
     Apply shift negative correction (SNC) on a graph built from a Pytorch model.
 
     Args:
         graph: Graph to apply SNC on.
         core_config: Quantization configuration.
+        fw_info: FrameworkInfo object with information about the specific framework's module.
 
     Returns:
         Graph after SNC.
@@ -229,6 +230,7 @@ def pytorch_apply_shift_negative_correction(graph: Graph,
     snc_node, linear_node, bypass_node, pad_node = shift_negative_activation_node_matchers()
     return apply_shift_negative_correction(graph,
                                            core_config,
+                                           fw_info,
                                            snc_node,
                                            linear_node,
                                            bypass_node,
@@ -240,5 +242,4 @@ def pytorch_apply_shift_negative_correction(graph: Graph,
                                            PADDING,
                                            BIAS,
                                            USE_BIAS,
-                                           get_activation_quantization_fn_factory,
                                            params_search_quantization_fn=params_search_quantization_fn)

model_compression_toolkit/core/pytorch/hessian/weights_hessian_scores_calculator_pytorch.py

@@ -23,6 +23,7 @@ from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS, MIN_HESS
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianScoresGranularity
 from model_compression_toolkit.core.pytorch.back2framework.float_model_builder import FloatPyTorchModelBuilder
+from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 from model_compression_toolkit.core.pytorch.hessian.hessian_scores_calculator_pytorch import \
     HessianScoresCalculatorPytorch
 from model_compression_toolkit.logger import Logger
@@ -91,14 +92,22 @@ class WeightsHessianScoresCalculatorPytorch(HessianScoresCalculatorPytorch):
             for i, ipt_node in enumerate(self.hessian_request.target_nodes):  # Per Interest point weights tensor
 
                 # Check if the target node's layer type is supported.
-                if not ipt_node.kernel_attr:
+                if not DEFAULT_PYTORCH_INFO.is_kernel_op(ipt_node.type):
                     Logger.critical(f"Hessian information with respect to weights is not supported for "
                                     f"{ipt_node.type} layers.")  # pragma: no cover
 
-                weights_tensor = getattr(getattr(model, ipt_node.name), ipt_node.kernel_attr)
+                # Get the weight attributes for the target node type
+                weights_attributes = DEFAULT_PYTORCH_INFO.get_kernel_op_attributes(ipt_node.type)
+
+                # Get the weight tensor for the target node
+                if len(weights_attributes) != 1:  # pragma: no cover
+                    Logger.critical(f"Currently, Hessian scores with respect to weights are supported only for nodes with a "
+                                    f"single weight attribute. {len(weights_attributes)} attributes found.")
+
+                weights_tensor = getattr(getattr(model, ipt_node.name), weights_attributes[0])
 
                 # Get the output channel index
-                output_channel_axis = ipt_node.channel_axis.output
+                output_channel_axis, _ = DEFAULT_PYTORCH_INFO.kernel_channels_mapping.get(ipt_node.type)
                 shape_channel_axis = [i for i in range(len(weights_tensor.shape))]
                 if self.hessian_request.granularity == HessianScoresGranularity.PER_OUTPUT_CHANNEL:
                     shape_channel_axis.remove(output_channel_axis)

model_compression_toolkit/core/pytorch/mixed_precision/configurable_activation_quantizer.py

@@ -20,7 +20,6 @@ from model_compression_toolkit.core.common.mixed_precision.configurable_quantize
     verify_candidates_descending_order, init_activation_quantizers
 from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
     CandidateNodeQuantizationConfig
-from model_compression_toolkit.core.pytorch.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.logger import Logger
 from mct_quantizers import QuantizationMethod
 from mct_quantizers import QuantizationTarget
@@ -68,7 +67,7 @@ class ConfigurableActivationQuantizer(BasePyTorchInferableQuantizer):
                 Logger.critical("Unsupported configuration: Mixing candidates with differing activation quantization states (enabled/disabled).")  # pragma: no cover
 
         # Setting layer's activation
-        self.activation_quantizers = init_activation_quantizers(self.node_q_cfg, get_activation_quantization_fn_factory)
+        self.activation_quantizers = init_activation_quantizers(self.node_q_cfg)
         self.active_quantization_config_index = max_candidate_idx  # initialize with first config as default
 
     def set_active_activation_quantizer(self, index: Optional[int]):