mct-nightly 2.4.0.20250630.629__py3-none-any.whl → 2.4.0.20250702.605__py3-none-any.whl

model_compression_toolkit/core/common/fusion/graph_fuser.py

@@ -14,12 +14,12 @@
 #  ==============================================================================
 
 import copy
-from typing import List, Tuple
+from typing import Tuple
 
 from model_compression_toolkit.core.common.fusion.fusing_info import FusingInfoGenerator
 from model_compression_toolkit.core.common.graph.base_graph import Graph, BaseNode, OutTensor
-from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import CandidateNodeQuantizationConfig
-from itertools import product
+from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
+    CandidateNodeQuantizationConfig, NodeQuantizationConfig
 
 
 class FusedLayerType:
@@ -30,6 +30,7 @@ class FusedLayerType:
     def __init__(self):
         self.__name__ = 'FusedLayer'
 
+
 class GraphFuser:
     def apply_node_fusion(self, graph: Graph) -> Graph:
         """
@@ -64,7 +65,6 @@ class GraphFuser:
 
         return graph_copy
 
-
     @staticmethod
     def _create_fused_node(fused_node_id: str, nodes: Tuple[BaseNode]) -> BaseNode:
         """
@@ -86,10 +86,15 @@ class GraphFuser:
                               weights={},
                               layer_class=FusedLayerType)
 
+        base_cfg = CandidateNodeQuantizationConfig(
+            activation_quantization_cfg=nodes[-1].quantization_cfg.base_quantization_cfg.activation_quantization_cfg,
+            weights_quantization_cfg=None
+        )
         activation_cfgs = [c.activation_quantization_cfg for c in nodes[-1].candidates_quantization_cfg]
-        fused_node.candidates_quantization_cfg = [
-            CandidateNodeQuantizationConfig(weights_quantization_cfg=None, activation_quantization_cfg=a) for a in
-            activation_cfgs]
+        candidates = [CandidateNodeQuantizationConfig(weights_quantization_cfg=None, activation_quantization_cfg=a)
+                      for a in activation_cfgs]
+        fused_node.quantization_cfg = NodeQuantizationConfig(base_quantization_cfg=base_cfg,
+                                                             candidates_quantization_cfg=candidates)
 
         # Keep the final configurations if they were set already.
         fused_node.final_weights_quantization_cfg = nodes[0].final_weights_quantization_cfg
@@ -158,5 +163,3 @@ class GraphFuser:
 
         # Finally, add the new fused node to the graph
         graph.add_node(fused_node)
-
-

model_compression_toolkit/core/common/graph/base_graph.py

@@ -39,6 +39,7 @@ from model_compression_toolkit.target_platform_capabilities.targetplatform2frame
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.framework_quantization_capabilities import \
     FrameworkQuantizationCapabilities
 
+
 def validate_graph_after_change(method: Callable) -> Callable:
     """
     Decorator for graph-mutating methods. After the decorated method executes,
@@ -120,28 +121,13 @@ class Graph(nx.MultiDiGraph, GraphSearches):
     def fusing_info(self, fusing_info: FusingInfo):
         self._fusing_info = fusing_info
 
-    def set_fqc(self,
-                fqc: FrameworkQuantizationCapabilities):
+    def set_fqc(self, fqc: FrameworkQuantizationCapabilities):
         """
         Set the graph's FQC.
         Args:
             fqc: FrameworkQuantizationCapabilities object.
         """
-        # validate graph nodes are either from the framework or a custom layer defined in the FQC
-        # Validate graph nodes are either built-in layers from the framework or custom layers defined in the FQC
-        fqc_layers = fqc.op_sets_to_layers.get_layers()
-        fqc_filtered_layers = [layer for layer in fqc_layers if isinstance(layer, LayerFilterParams)]
-        for n in self.nodes:
-            is_node_in_fqc = any([n.is_match_type(_type) for _type in fqc_layers]) or \
-                             any([n.is_match_filter_params(filtered_layer) for filtered_layer in fqc_filtered_layers])
-            if n.is_custom:
-                if not is_node_in_fqc:
-                    Logger.critical(f'MCT does not support optimizing Keras custom layers. Found a layer of type {n.type}. '
-                                    ' Please add the custom layer to Framework Quantization Capabilities (FQC), or file a feature '
-                                    'request or an issue if you believe this should be supported.')  # pragma: no cover
-                if any([qc.default_weight_attr_config.enable_weights_quantization for qc in n.get_qco(fqc).quantization_configurations]):
-                    Logger.critical(f'Layer identified: {n.type}. MCT does not support weight quantization for Keras custom layers.')  # pragma: no cover
-
+        # TODO irena: this is only passed for negative shift activation.
         self.fqc = fqc
 
     def get_topo_sorted_nodes(self):
@@ -578,7 +564,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
             A list of nodes that their weights can be configured (namely, has one or more weight qc candidate).
         """
         # configurability is only relevant for kernel attribute quantization
-        potential_conf_nodes = [n for n in list(self) if n.is_kernel_op]
+        potential_conf_nodes = [n for n in self.nodes if n.kernel_attr]
 
         def is_configurable(n):
             return n.is_configurable_weight(n.kernel_attr) and (not n.reuse or include_reused_nodes)
@@ -693,10 +679,8 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         """
         Gets the final number of bits for quantization of each weights' configurable layer.
 
-        Args:
-            fw_info: fw_info: FrameworkInfo object with information about the specific framework's model.
-
-        Returns: A list of pairs of (node type, node's weights quantization bitwidth).
+        Returns:
+            A list of pairs of (node type, node's weights quantization bitwidth).
 
         """
         sorted_conf_weights = self.get_sorted_weights_configurable_nodes()
@@ -876,32 +860,31 @@ class Graph(nx.MultiDiGraph, GraphSearches):
 
         return intermediate_nodes, next_node
 
+    # TODO irena move to load_fqc and clean up tests (currently tests_pytest/common_tests/unit_tests/core/graph/test_base_graph.py)
     def override_fused_node_activation_quantization_candidates(self):
         """
         Override fused node activation quantization candidates for all nodes in fused operations,
         except for the last node in each fused group.
         Update the value of quantization_config with the value of op_quaitization_cfg from FusingInfo.
         """
-        from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import CandidateNodeQuantizationConfig
-
         nodes_in_fln = self.fusing_info.get_inner_fln_nodes()
         for node in nodes_in_fln:
             fused_node_op_id = self.fusing_info.get_fused_op_id_for_node(node.name)
-            fusiong_op_quaitization_cfg = self.fusing_info.get_fused_op_quantization_config(fused_node_op_id)             
-            org_candidate = node.candidates_quantization_cfg[0]
-            if fusiong_op_quaitization_cfg is not None and fusiong_op_quaitization_cfg.enable_activation_quantization:
-                # Set ActivationQuantizationMode to FLN_QUANT and update the value of quantization_config
-                activation_quantization_cfg = NodeActivationQuantizationConfig(qc=org_candidate,
-                                                                               op_cfg=fusiong_op_quaitization_cfg,
-                                                                               activation_quantization_fn=org_candidate.activation_quantization_cfg.activation_quantization_fn,
-                                                                               activation_quantization_params_fn=org_candidate.activation_quantization_cfg.activation_quantization_params_fn)
-                activation_quantization_cfg.quant_mode = ActivationQuantizationMode.FLN_QUANT
-                for qc in node.candidates_quantization_cfg:
-                    qc.activation_quantization_cfg = activation_quantization_cfg
+            fusing_op_quantization_cfg = self.fusing_info.get_fused_op_quantization_config(fused_node_op_id)
+            if fusing_op_quantization_cfg is not None and fusing_op_quantization_cfg.enable_activation_quantization:
+                def update(qc):
+                    qc.activation_quantization_cfg = NodeActivationQuantizationConfig(fusing_op_quantization_cfg)
+                    qc.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.FLN_QUANT
+                node.quantization_cfg.update_all(update, remove_duplicates=True)
             else:
-                # Set ActivationQuantizationMode to FLN_NO_QUANT
+                node.quantization_cfg.update_activation_quantization_mode(ActivationQuantizationMode.FLN_NO_QUANT)
+                # Remove duplicate candidates. We cannot compare whole candidates since activation configs might not
+                # be identical, but we do want to treat them as such. So we only check duplication by weight configs.
+                uniq_qcs = []
                 for qc in node.candidates_quantization_cfg:
-                    qc.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.FLN_NO_QUANT
+                    if not any(qc.weights_quantization_cfg == uqc.weights_quantization_cfg for uqc in uniq_qcs):
+                        uniq_qcs.append(qc)
+                node.quantization_cfg.candidates_quantization_cfg = uniq_qcs
 
     def validate(self):
         """

model_compression_toolkit/core/common/graph/base_node.py

@@ -21,15 +21,11 @@ import numpy as np
 from model_compression_toolkit.core.common.framework_info import get_fw_info, ChannelAxisMapping
 from model_compression_toolkit.constants import WEIGHTS_NBITS_ATTRIBUTE, CORRECTED_BIAS_ATTRIBUTE, \
     ACTIVATION_N_BITS_ATTRIBUTE, FP32_BYTES_PER_PARAMETER
+from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import NodeQuantizationConfig
 from model_compression_toolkit.core.common.quantization.node_quantization_config import WeightsAttrQuantizationConfig, \
     ActivationQuantizationMode
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import QuantizationConfigOptions, \
-    OpQuantizationConfig
-from model_compression_toolkit.target_platform_capabilities.schema.schema_functions import max_input_activation_n_bits
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework import LayerFilterParams
-from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.framework_quantization_capabilities import \
-    FrameworkQuantizationCapabilities
 
 
 WeightAttrT = Union[str, int]
@@ -43,7 +39,6 @@ class NodeFrameworkInfo(NamedTuple):
     out_channel_axis: int
     minmax: Tuple[float, float]
     kernel_attr: str
-    is_kernel_op: bool
 
 
 class BaseNode:
@@ -95,7 +90,7 @@ class BaseNode:
         self.inputs_as_list = inputs_as_list
         self.final_weights_quantization_cfg = None
         self.final_activation_quantization_cfg = None
-        self.candidates_quantization_cfg = None
+        self.quantization_cfg: NodeQuantizationConfig = None
         self.prior_info = None
         self.has_activation = has_activation
         self.is_custom = is_custom
@@ -108,7 +103,6 @@ class BaseNode:
             fw_info.get_out_channel_axis(node_type),
             fw_info.get_layer_min_max(node_type, framework_attr),
             fw_info.get_kernel_op_attribute(node_type),
-            fw_info.is_kernel_op(node_type)
         )
 
     def _assert_fw_info_exists(self):
@@ -162,15 +156,9 @@ class BaseNode:
         return self.node_fw_info.kernel_attr
 
     @property
-    def is_kernel_op(self) -> bool:
-        """
-        Check if kernel exists for the node.
-
-        Returns:
-            Whether the node has a kernel or not.
-        """
-        self._assert_fw_info_exists()
-        return self.node_fw_info.is_kernel_op
+    def candidates_quantization_cfg(self):
+        assert self.quantization_cfg
+        return self.quantization_cfg.candidates_quantization_cfg
 
     @property
     def type(self):
@@ -181,15 +169,6 @@ class BaseNode:
         """
         return self.layer_class
 
-    def get_has_activation(self):
-        """
-        Returns has_activation attribute.
-
-        Returns: Whether the node has activation to quantize.
-
-        """
-        return self.has_activation
-
     @property
     def has_positional_weights(self):
         """
@@ -646,8 +625,9 @@ class BaseNode:
         Returns: True if the node has at list one quantization configuration candidate with activation quantization enabled.
         """
 
-        return len(self.candidates_quantization_cfg) > 0 and \
-            any([c.activation_quantization_cfg.enable_activation_quantization for c in self.candidates_quantization_cfg])
+        return (len(self.candidates_quantization_cfg) > 0 and
+                any([c.activation_quantization_cfg.enable_activation_quantization
+                     for c in self.candidates_quantization_cfg]))
 
     def get_all_weights_attr_candidates(self, attr: str) -> List[WeightsAttrQuantizationConfig]:
         """
@@ -663,79 +643,6 @@ class BaseNode:
         # the inner method would log an exception.
         return [c.weights_quantization_cfg.get_attr_config(attr) for c in self.candidates_quantization_cfg]
 
-    def get_qco(self, fqc: FrameworkQuantizationCapabilities) -> QuantizationConfigOptions:
-        """
-        Get the QuantizationConfigOptions of the node according
-        to the mappings from layers/LayerFilterParams to the OperatorsSet in the TargetPlatformCapabilities.
-
-        Args:
-            fqc: FQC to extract the QuantizationConfigOptions for the node.
-
-        Returns:
-            QuantizationConfigOptions of the node.
-        """
-
-        if fqc is None:
-            Logger.critical(f'Can not retrieve QC options for None FQC')  # pragma: no cover
-
-        for fl, qco in fqc.filterlayer2qco.items():
-            if self.is_match_filter_params(fl):
-                return qco
-        # Extract qco with is_match_type to overcome mismatch of function types in TF 2.15
-        matching_qcos = [_qco for _type, _qco in fqc.layer2qco.items() if self.is_match_type(_type)]
-        if matching_qcos:
-            if all([_qco == matching_qcos[0] for _qco in matching_qcos]):
-                return matching_qcos[0]
-            else:
-                Logger.critical(f"Found duplicate qco types for node '{self.name}' of type '{self.type}'!")  # pragma: no cover
-        return fqc.tpc.default_qco
-
-    def filter_node_qco_by_graph(self, fqc: FrameworkQuantizationCapabilities,
-                                 next_nodes: List, node_qc_options: QuantizationConfigOptions
-                                 ) -> Tuple[OpQuantizationConfig, List[OpQuantizationConfig]]:
-        """
-        Filter quantization config options that don't match the graph.
-        A node may have several quantization config options with 'activation_n_bits' values, and
-        the next nodes in the graph may support different bit-width as input activation. This function
-        filters out quantization config that don't comply to these attributes.
-
-        Args:
-            fqc: FQC to extract the QuantizationConfigOptions for the next nodes.
-            next_nodes: Output nodes of current node.
-            node_qc_options: Node's QuantizationConfigOptions.
-
-        Returns:
-
-        """
-        # Filter quantization config options that don't match the graph.
-        _base_config = node_qc_options.base_config
-        _node_qc_options = node_qc_options.quantization_configurations
-        if len(next_nodes):
-            next_nodes_qc_options = [_node.get_qco(fqc) for _node in next_nodes]
-            next_nodes_supported_input_bitwidth = min([max_input_activation_n_bits(op_cfg)
-                                                       for qc_opts in next_nodes_qc_options
-                                                       for op_cfg in qc_opts.quantization_configurations])
-
-            # Filter node's QC options that match next nodes input bit-width.
-            _node_qc_options = [_option for _option in _node_qc_options
-                                if _option.activation_n_bits <= next_nodes_supported_input_bitwidth]
-            if len(_node_qc_options) == 0:
-                Logger.critical(f"Graph doesn't match FQC bit configurations: {self} -> {next_nodes}.")  # pragma: no cover
-
-            # Verify base config match
-            if any([node_qc_options.base_config.activation_n_bits > max_input_activation_n_bits(qc_opt.base_config)
-                    for qc_opt in next_nodes_qc_options]):
-                # base_config activation bits doesn't match next node supported input bit-width -> replace with
-                # a qco from quantization_configurations with maximum activation bit-width.
-                if len(_node_qc_options) > 0:
-                    output_act_bitwidth = {qco.activation_n_bits: i for i, qco in enumerate(_node_qc_options)}
-                    _base_config = _node_qc_options[output_act_bitwidth[max(output_act_bitwidth)]]
-                    Logger.warning(f"Node {self} base quantization config changed to match Graph and FQC configuration.\nCause: {self} -> {next_nodes}.")
-                else:
-                    Logger.critical(f"Graph doesn't match FQC bit configurations: {self} -> {next_nodes}.")  # pragma: no cover
-
-        return _base_config, _node_qc_options
-
     def is_match_type(self, _type: Type) -> bool:
         """
         Check if input type matches the node type, either in instance type or in type name.
@@ -768,7 +675,7 @@ class BaseNode:
             return False
 
         # Get attributes from node to filter
-        layer_config = self.framework_attr
+        layer_config = self.framework_attr.copy()
         if hasattr(self, "op_call_kwargs"):
             layer_config.update(self.op_call_kwargs)
 
@@ -812,11 +719,11 @@ class BaseNode:
             the candidates in descending order.
         The operation is done inplace.
         """
-        if self.candidates_quantization_cfg is not None:
+        if self.quantization_cfg.candidates_quantization_cfg is not None:
             if self.kernel_attr is not None:
-                self.candidates_quantization_cfg.sort(
+                self.quantization_cfg.candidates_quantization_cfg.sort(
                     key=lambda c: (c.weights_quantization_cfg.get_attr_config(self.kernel_attr).weights_n_bits,
                                    c.activation_quantization_cfg.activation_n_bits), reverse=True)
             else:
-                self.candidates_quantization_cfg.sort(key=lambda c: c.activation_quantization_cfg.activation_n_bits,
-                                                      reverse=True)
+                self.quantization_cfg.candidates_quantization_cfg.sort(
+                    key=lambda c: c.activation_quantization_cfg.activation_n_bits, reverse=True)

model_compression_toolkit/core/common/graph/functional_node.py

@@ -103,4 +103,4 @@ class FunctionalNode(BaseNode):
 
         """
         names_match = _type.__name__ == self.type.__name__
-        return super().is_match_type(_type) or names_match
+        return names_match or super().is_match_type(_type)

model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py

@@ -19,9 +19,8 @@ from model_compression_toolkit.constants import VIRTUAL_ACTIVATION_WEIGHTS_NODE_
     VIRTUAL_WEIGHTS_SUFFIX, VIRTUAL_ACTIVATION_SUFFIX, FLOAT_BITWIDTH
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
-    CandidateNodeQuantizationConfig
+    CandidateNodeQuantizationConfig, NodeQuantizationConfig
 from model_compression_toolkit.core.common.quantization.node_quantization_config import ActivationQuantizationMode
-from model_compression_toolkit.core.common.framework_info import DEFAULT_KERNEL_ATTRIBUTE
 
 
 class VirtualNode(BaseNode, abc.ABC):
@@ -76,8 +75,11 @@ class VirtualSplitWeightsNode(VirtualSplitNode):
 
         self.name = origin_node.name + VIRTUAL_WEIGHTS_SUFFIX
 
-        self.candidates_quantization_cfg = origin_node.get_unique_weights_candidates(kernel_attr)
-        for c in self.candidates_quantization_cfg:
+        self.quantization_cfg = NodeQuantizationConfig(
+            candidates_quantization_cfg=origin_node.get_unique_weights_candidates(kernel_attr),
+            base_quantization_cfg=None, validate=False
+        )
+        for c in self.quantization_cfg.candidates_quantization_cfg:
             c.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.NO_QUANT
             c.activation_quantization_cfg.activation_n_bits = FLOAT_BITWIDTH
 
@@ -106,10 +108,9 @@ class VirtualSplitActivationNode(VirtualSplitNode):
         self.weights = {}
         self.layer_class = activation_class
 
-        self.candidates_quantization_cfg = origin_node.get_unique_activation_candidates()
-        for c in self.candidates_quantization_cfg:
-            c.weights_quantization_cfg.enable_weights_quantization = False
-            c.weights_quantization_cfg.weights_n_bits = FLOAT_BITWIDTH
+        self.quantization_cfg = NodeQuantizationConfig(candidates_quantization_cfg=origin_node.get_unique_activation_candidates(),
+                                                       base_quantization_cfg=None, validate=False)
+        self.quantization_cfg.disable_weights_quantization()
 
 
 class VirtualActivationWeightsNode(VirtualNode):
@@ -143,7 +144,7 @@ class VirtualActivationWeightsNode(VirtualNode):
         weights = weights_node.weights.copy()
         act_node_w_rename = {}
         if act_node.weights:
-            if act_node.kernel_attr != DEFAULT_KERNEL_ATTRIBUTE:
+            if act_node.kernel_attr:
                 raise NotImplementedError(f'Node {act_node} with kernel cannot be used as activation for '
                                           f'VirtualActivationWeightsNode.')
             if act_node.has_any_configurable_weight():
@@ -200,4 +201,5 @@ class VirtualActivationWeightsNode(VirtualNode):
         v_candidates.sort(key=lambda c: (c.weights_quantization_cfg.get_attr_config(weights_node.kernel_attr).weights_n_bits,
                                          c.activation_quantization_cfg.activation_n_bits), reverse=True)
 
-        self.candidates_quantization_cfg = v_candidates
+        self.quantization_cfg = NodeQuantizationConfig(candidates_quantization_cfg=v_candidates,
+                                                       base_quantization_cfg=None, validate=False)

model_compression_toolkit/core/common/mixed_precision/configurable_quantizer_utils.py

@@ -18,6 +18,8 @@ import numpy as np
 
 from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
     CandidateNodeQuantizationConfig
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import (get_activation_quantization_fn,
+                                                                                          get_weights_quantization_fn)
 
 
 def verify_candidates_descending_order(node_q_cfg: List[CandidateNodeQuantizationConfig],
@@ -77,20 +79,21 @@ def init_quantized_weights(node_q_cfg: List[CandidateNodeQuantizationConfig],
     quantized_weights = []
     for qc in node_q_cfg:
         qc_weights_attr = qc.weights_quantization_cfg.get_attr_config(kernel_attr)
-        q_weight = qc_weights_attr.weights_quantization_fn(float_weights,
-                                                           qc_weights_attr.weights_n_bits,
-                                                           True,
-                                                           qc_weights_attr.weights_quantization_params,
-                                                           qc_weights_attr.weights_per_channel_threshold,
-                                                           qc_weights_attr.weights_channels_axis[
-                                                               0])  # output channel axis
+        weights_quantization_fn = get_weights_quantization_fn(qc_weights_attr.weights_quantization_method)
+        q_weight = weights_quantization_fn(float_weights,
+                                           qc_weights_attr.weights_n_bits,
+                                           True,
+                                           qc_weights_attr.weights_quantization_params,
+                                           qc_weights_attr.weights_per_channel_threshold,
+                                           qc_weights_attr.weights_channels_axis[0])  # output channel axis
 
         quantized_weights.append(fw_tensor_convert_func(q_weight))
 
     return quantized_weights
 
 
-def init_activation_quantizers(node_q_cfg: List[CandidateNodeQuantizationConfig]) -> List:
+def init_activation_quantizers(node_q_cfg: List[CandidateNodeQuantizationConfig],
+                               get_activation_quantization_fn_factory: Callable) -> List:
     """
     Builds a list of quantizers for each of the bitwidth candidates for activation quantization,
     to be stored and used during MP search.
@@ -98,6 +101,7 @@ def init_activation_quantizers(node_q_cfg: List[CandidateNodeQuantizationConfig]
     Args:
         node_q_cfg: Quantization configuration candidates of the node that generated the layer that will
                     use this quantizer.
+        get_activation_quantization_fn_factory: activation quantization functions factory.
 
     Returns: a list of activation quantizers - for each bitwidth and layer's attribute to be quantized.
     """
@@ -105,6 +109,7 @@ def init_activation_quantizers(node_q_cfg: List[CandidateNodeQuantizationConfig]
     activation_quantizers = []
     for index, qc in enumerate(node_q_cfg):
         q_activation = node_q_cfg[index].activation_quantization_cfg
-        activation_quantizers.append(q_activation.quantize_node_output)
+        quantizer = get_activation_quantization_fn(q_activation, get_activation_quantization_fn_factory)
+        activation_quantizers.append(quantizer)
 
     return activation_quantizers

model_compression_toolkit/core/common/mixed_precision/mixed_precision_candidates_filter.py

@@ -12,17 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import numpy as np
-
-from model_compression_toolkit.core import ResourceUtilization, FrameworkInfo
+from model_compression_toolkit.core import ResourceUtilization
 from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.framework_quantization_capabilities import \
-    FrameworkQuantizationCapabilities
 
 
 def filter_candidates_for_mixed_precision(graph: Graph,
-                                          target_resource_utilization: ResourceUtilization,
-                                          fqc: FrameworkQuantizationCapabilities):
+                                          target_resource_utilization: ResourceUtilization):
     """
     Filters out candidates in case of mixed precision search for only weights or activation compression.
     For instance, if running only weights compression - filters out candidates of activation configurable nodes
@@ -34,8 +29,6 @@ def filter_candidates_for_mixed_precision(graph: Graph,
     Args:
         graph: A graph representation of the model to be quantized.
         target_resource_utilization: The resource utilization of the target device.
-        fqc: FrameworkQuantizationCapabilities object that describes the desired inference target platform.
-
     """
 
     tru = target_resource_utilization
@@ -47,20 +40,21 @@ def filter_candidates_for_mixed_precision(graph: Graph,
         # filter out candidates activation only configurable node
         activation_configurable_nodes = [n for n in graph.get_activation_configurable_nodes()]
         for n in activation_configurable_nodes:
-            base_cfg_nbits = n.get_qco(fqc).base_config.activation_n_bits
-            filtered_conf = [c for c in n.candidates_quantization_cfg if
+            base_cfg_nbits = n.quantization_cfg.base_quantization_cfg.activation_quantization_cfg.activation_n_bits
+            filtered_cfgs = [c for c in n.candidates_quantization_cfg if
                              c.activation_quantization_cfg.enable_activation_quantization and
                              c.activation_quantization_cfg.activation_n_bits == base_cfg_nbits]
 
-            n.candidates_quantization_cfg = filtered_conf
+            n.quantization_cfg.candidates_quantization_cfg = filtered_cfgs
 
     elif tru.activation_restricted() and not tru.weight_restricted():
         # Running mixed precision for activation compression only -
         # filter out candidates weights only configurable node
         weight_configurable_nodes = [n for n in graph.get_weights_configurable_nodes()]
         for n in weight_configurable_nodes:
-            base_cfg_nbits = n.get_qco(fqc).base_config.attr_weights_configs_mapping[n.kernel_attr].weights_n_bits
-            filtered_conf = [c for c in n.candidates_quantization_cfg if
+            base_cfg_nbits = (n.quantization_cfg.base_quantization_cfg.weights_quantization_cfg.
+                              get_attr_config(n.kernel_attr).weights_n_bits)
+            filtered_cfgs = [c for c in n.candidates_quantization_cfg if
                              c.weights_quantization_cfg.get_attr_config(n.kernel_attr).enable_weights_quantization and
                              c.weights_quantization_cfg.get_attr_config(n.kernel_attr).weights_n_bits == base_cfg_nbits]
-            n.candidates_quantization_cfg = filtered_conf
+            n.quantization_cfg.candidates_quantization_cfg = filtered_cfgs

model_compression_toolkit/core/common/mixed_precision/sensitivity_eval/metric_calculators.py

@@ -392,9 +392,8 @@ class DistanceMetricCalculator(MetricCalculator):
         """
 
         return [n.node for n in graph.get_outputs()
-                if (n.node.is_kernel_op and
-                    n.node.is_weights_quantization_enabled(n.node.kernel_attr)) or
-                n.node.is_activation_quantization_enabled()]
+                if (n.node.kernel_attr and n.node.is_weights_quantization_enabled(n.node.kernel_attr))
+                or n.node.is_activation_quantization_enabled()]
 
     @staticmethod
     def bound_num_interest_points(sorted_ip_list: List[BaseNode], num_ip_factor: float) -> List[BaseNode]:

model_compression_toolkit/core/common/network_editors/__init__.py

@@ -13,7 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.core.common.network_editors.actions import ChangeCandidatesWeightsQuantConfigAttr, ChangeFinalWeightsQuantConfigAttr, ChangeCandidatesActivationQuantConfigAttr, ChangeQuantizationParamFunction, ChangeCandidatesActivationQuantizationMethod, ChangeFinalWeightsQuantizationMethod, ChangeCandidatesWeightsQuantizationMethod, ChangeFinalActivationQuantConfigAttr
+from model_compression_toolkit.core.common.network_editors.actions import (
+    ChangeCandidatesWeightsQuantConfigAttr,
+    ChangeFinalWeightsQuantConfigAttr,
+    ChangeCandidatesActivationQuantConfigAttr,
+    ChangeCandidatesActivationQuantizationMethod,
+    ChangeFinalWeightsQuantizationMethod,
+    ChangeCandidatesWeightsQuantizationMethod,
+    ChangeFinalActivationQuantConfigAttr)
 from model_compression_toolkit.core.common.network_editors.actions import EditRule
 from model_compression_toolkit.core.common.network_editors.node_filters import NodeTypeFilter, NodeNameScopeFilter, \
     NodeNameFilter