mct-nightly 2.4.0.20250616.616__py3-none-any.whl → 2.4.0.20250618.606__py3-none-any.whl

model_compression_toolkit/core/pytorch/back2framework/mixed_precision_model_builder.py

@@ -23,7 +23,6 @@ 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 \
@@ -38,14 +37,12 @@ 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.
         """
 
@@ -53,7 +50,6 @@ 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)
@@ -77,17 +73,16 @@ class MixedPrecisionPyTorchModelBuilder(PyTorchModelBuilder):
             ValueError: if kernel attribute is quantized but not configurable.
         """
 
-        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):
+        if n.kernel_attr is None or not n.is_weights_quantization_enabled(n.kernel_attr):
             return layer
-        if not n.is_configurable_weight(kernel_attr):  # pragma: no cover
+        if not n.is_configurable_weight(n.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={
-                                              kernel_attr: ConfigurableWeightsQuantizer(
+                                              n.kernel_attr: ConfigurableWeightsQuantizer(
                                                   **self._get_weights_configurable_quantizer_kwargs(n,
-                                                                                                    kernel_attr),
-                                                  kernel_attr=kernel_attr)})
+                                                                                                    n.kernel_attr),
+                                                  kernel_attr=n.kernel_attr)})
 
     def _get_weights_configurable_quantizer_kwargs(self, n: BaseNode, attr: str) -> Dict[str, Any]:
         """
@@ -147,14 +142,13 @@ 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(self.fw_info)
+        n.sort_node_candidates()
 
         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': kernel_attr})] \
+                                                                    'kernel_attr': n.kernel_attr})] \
                                 * num_of_outputs
 
         # Holder by definition uses a single quantizer for the activation quantization
@@ -177,7 +171,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(self.fw_info)}
+                            for n in self.graph.get_configurable_sorted_nodes()}
 
         return model, user_info, conf_node2layers
 
@@ -230,8 +224,7 @@ class MixedPrecisionPyTorchModelBuilder(PyTorchModelBuilder):
 
         """
         # Only layers with kernel op are considered weights configurable
-        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)
+        weights_quant = False if n.kernel_attr is None else n.is_weights_quantization_enabled(n.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,7 +30,6 @@ 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
@@ -364,7 +363,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(DEFAULT_PYTORCH_INFO)
+        configurable_nodes = self.graph.get_configurable_sorted_nodes_names()
         for node in self.node_sort:
             op_func = self._get_op_func(node, configurable_nodes)
             input_tensors = _build_input_tensors_list(node,
@@ -440,7 +439,6 @@ 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):
@@ -449,7 +447,6 @@ 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.
@@ -457,7 +454,6 @@ 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,7 +21,6 @@ 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
 
 
@@ -93,7 +92,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.
 
@@ -104,7 +103,7 @@ class QuantizedLayerWrapper(torch.nn.Module):
             None.
         """
 
-        self.weight_attrs = DEFAULT_PYTORCH_INFO.get_kernel_op_attributes(n.type)
+        self.weight_attrs = [n.kernel_attr]
 
         # 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

@@ -23,7 +23,6 @@ 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, \
     PytorchModel
-from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
 
 
 class QuantizedPyTorchModel(PytorchModel):
@@ -70,20 +69,17 @@ 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,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from torch.nn import Hardsigmoid, ReLU, ReLU6, Softmax, Sigmoid, GELU, SELU
-from torch.nn.functional import hardsigmoid, relu, relu6, softmax, gelu, selu
+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 Conv2d, ConvTranspose2d, Linear
 from torch import sigmoid
 
-from model_compression_toolkit.defaultdict import DefaultDict
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo, DEFAULT_KERNEL_ATTRIBUTES
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo, set_fw_info, ChannelAxisMapping
 from mct_quantizers import QuantizationMethod
 from model_compression_toolkit.constants import SOFTMAX_THRESHOLD
 from model_compression_toolkit.core.pytorch.constants import KERNEL
@@ -26,73 +28,97 @@ from model_compression_toolkit.core.pytorch.quantizer.fake_quant_builder import
     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)
+
+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),
+                              }
+
+    """
+    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}
+
+    @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, 1)
+
+
+def set_pytorch_info(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        set_fw_info(PyTorchInfo)
+        return func(*args, **kwargs)
+    return wrapper

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

@@ -21,19 +21,18 @@ 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 import FrameworkInfo
+from model_compression_toolkit.core.common.framework_info import get_fw_info
 
 
 class FunctionalConvSubstitution(common.BaseSubstitution):
     """
     Substitute functional convolutions with Layers
     """
-    def __init__(self, fw_info: FrameworkInfo):
+    def __init__(self):
         """
         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,
@@ -56,7 +55,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 = self.fw_info.kernel_channels_mapping.get(new_layer)
+        out_channel_index, in_channel_index = get_fw_info().get_kernel_channels(new_layer)
 
         # Create new node of layer convolution
         if 1 not in func_node.weights:

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

@@ -46,17 +46,15 @@ class ScaleEqualization(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig,
-                 fw_info: FrameworkInfo):
+                 quant_config: QuantizationConfig):
         """
         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, fw_info=fw_info, matcher_instance=MATCHER,
+        super().__init__(quant_config=quant_config, matcher_instance=MATCHER,
                          kernel_str=KERNEL, bias_str=BIAS)
 
 
@@ -66,15 +64,13 @@ class ScaleEqualizationWithPad(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig,
-                 fw_info: FrameworkInfo):
+                 quant_config: QuantizationConfig):
         """
         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, fw_info=fw_info, matcher_instance=MATCHER_WITH_PAD,
+        super().__init__(quant_config=quant_config, matcher_instance=MATCHER_WITH_PAD,
                          kernel_str=KERNEL, bias_str=BIAS)

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

@@ -214,15 +214,13 @@ def is_padding_node_and_node_has_padding(pad_node_to_consider: BaseNode,
 
 
 def pytorch_apply_shift_negative_correction(graph: Graph,
-                                            core_config: CoreConfig,
-                                            fw_info: FrameworkInfo) -> Graph:
+                                            core_config: CoreConfig) -> 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.
@@ -230,7 +228,6 @@ 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,

model_compression_toolkit/core/pytorch/hessian/weights_hessian_scores_calculator_pytorch.py

@@ -23,7 +23,6 @@ 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
@@ -92,22 +91,14 @@ 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 DEFAULT_PYTORCH_INFO.is_kernel_op(ipt_node.type):
+                if not ipt_node.is_kernel_op:
                     Logger.critical(f"Hessian information with respect to weights is not supported for "
                                     f"{ipt_node.type} layers.")  # pragma: no cover
 
-                # 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])
+                weights_tensor = getattr(getattr(model, ipt_node.name), ipt_node.kernel_attr)
 
                 # Get the output channel index
-                output_channel_axis, _ = DEFAULT_PYTORCH_INFO.kernel_channels_mapping.get(ipt_node.type)
+                output_channel_axis = ipt_node.channel_axis.output
                 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)