Initial quantization preparation (#1475)

Author: irenaby

model_compression_toolkit/core/common/back2framework/base_model_builder.py

@@ -15,7 +15,6 @@
 from abc import ABC, abstractmethod
 from typing import Any, Tuple
 
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.user_info import UserInformation
 

model_compression_toolkit/core/common/framework_info.py

@@ -16,18 +16,12 @@
 
 from collections.abc import Callable
 from enum import Enum
-from typing import Dict, Any, Tuple, NamedTuple
+from typing import Dict, Any, Tuple, NamedTuple, Optional
 from abc import ABC, abstractmethod
 
 from mct_quantizers import QuantizationMethod
 
 
-# Default value to use for ops without kernel.
-# This is a weird default, but it's used all over the place, so for now only extract it to const so that it can be
-# referenced by variable instead of hard-coded.
-DEFAULT_KERNEL_ATTRIBUTE = None
-
-
 class ChannelAxis(Enum):
     """
 
@@ -63,7 +57,6 @@ class FrameworkInfo(ABC):
         kernel_ops_attribute_mapping (Dict): Dictionary from a framework operator to its weight attribute to quantize.
         out_channel_axis_mapping (Dict): Dictionary of output channels of the model's layers (for computing statistics per-channel).
         _layer_min_max_mapping (Dict[Any, tuple]): Dictionary from a layer to its min/max output values.
-
     """
 
     activation_quantizer_mapping: Dict[QuantizationMethod, Callable]
@@ -75,7 +68,7 @@ class FrameworkInfo(ABC):
     _default_channel_mapping = ChannelAxisMapping(None, None)
 
     @classmethod
-    def get_kernel_op_attribute(cls, node_type: Any) -> str:
+    def get_kernel_op_attribute(cls, node_type: Any) -> Optional[str]:
         """
         Get attribute of a layer's weight to quantize.
 
@@ -85,20 +78,7 @@ def get_kernel_op_attribute(cls, node_type: Any) -> str:
         Returns:
             Attribute the layer has and should be quantized.
         """
-        return cls.kernel_ops_attribute_mapping.get(node_type, DEFAULT_KERNEL_ATTRIBUTE)
-
-    @classmethod
-    def is_kernel_op(cls, node_type: Any) -> bool:
-        """
-        Check is the node is a kernel operation.
-
-        Args:
-            node_type: Layer to get its attributes.
-
-        Returns:
-            True if node type is a kernel operation, else False.
-        """
-        return node_type in cls.kernel_ops_attribute_mapping
+        return cls.kernel_ops_attribute_mapping.get(node_type)
 
     @classmethod
     def get_layer_min_max(cls, layer: Any, fw_attrs: Dict) -> Tuple[float, float]:
@@ -169,7 +149,6 @@ def get_fw_info():
     Returns: FrameworkInfo class.
     """
     assert _current_framework_info is not None, "fw_info isn't initialized."
-    assert issubclass(_current_framework_info, FrameworkInfo), "fw_info isn't initialized to a FrameworkInfo class."
     return _current_framework_info
 
 

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 @@ def apply_node_fusion(self, graph: Graph) -> Graph:
 
         return graph_copy
 
-
     @staticmethod
     def _create_fused_node(fused_node_id: str, nodes: Tuple[BaseNode]) -> BaseNode:
         """
@@ -86,10 +86,15 @@ def _create_fused_node(fused_node_id: str, nodes: Tuple[BaseNode]) -> BaseNode:
                               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 @@ def _replace_nodes_with_fused_node(graph: Graph,
 
         # 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.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 @@ def fusing_info(self) -> FusingInfo:
     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 @@ def get_weights_configurable_nodes(self,
             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 @@ def get_final_weights_config(self) -> List[Tuple[BaseNode, int]]:
         """
         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,36 @@ def _find_intermediate_and_exit_nodes(self, entry_node: BaseNode, fw_impl: Any)
 
         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.activation_quantization_fn,
+                        qc.activation_quantization_cfg.activation_quantization_params_fn
+                    )
+                    qc.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.FLN_QUANT
+                node.quantization_cfg.update_all(update)
+                node.quantization_cfg.remove_duplicates()
             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):
         """