Add an option to dynamically disable configurable quantizers, separate configuration api for activation and weights quantizers (#1438)

* add an option to disable configurable quantizers
* separate configuring activation and weight quantizers

Author: irenaby

model_compression_toolkit/core/common/framework_implementation.py

@@ -13,31 +13,31 @@
 # limitations under the License.
 # ==============================================================================
 from abc import ABC, abstractmethod
-from typing import Callable, Any, List, Tuple, Dict, Generator
+from typing import Callable, Any, List, Tuple, Generator, Type
 
 import numpy as np
 
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
-from model_compression_toolkit.core import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
-from model_compression_toolkit.core.common.collectors.statistics_collector import BaseStatsCollector
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.graph.base_graph import Graph
-from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianInfoService
-from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
 from model_compression_toolkit.core.common.node_prior_info import NodePriorInfo
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
-from model_compression_toolkit.core.common.user_info import UserInformation
 
 
 class FrameworkImplementation(ABC):
     """
     An abstract class with abstract methods that should be implemented when supporting a new
     framework in MCT.
     """
+    weights_quant_layer_cls: Type
+    activation_quant_layer_cls: Type
+    configurable_weights_quantizer_cls: Type
+    configurable_activation_quantizer_cls: Type
 
     @property
     def constants(self):
@@ -327,33 +327,6 @@ def get_substitutions_after_second_moment_correction(self, quant_config: Quantiz
                              f'framework\'s get_substitutions_after_second_moment_correction '
                              f'method.')  # pragma: no cover
 
-    @abstractmethod
-    def get_sensitivity_evaluator(self,
-                                  graph: Graph,
-                                  quant_config: MixedPrecisionQuantizationConfig,
-                                  representative_data_gen: Callable,
-                                  fw_info: FrameworkInfo,
-                                  hessian_info_service: HessianInfoService = None,
-                                  disable_activation_for_metric: bool = False) -> SensitivityEvaluation:
-        """
-        Creates and returns an object which handles the computation of a sensitivity metric for a mixed-precision
-        configuration (comparing to the float model).
-
-        Args:
-            graph: Graph to build its float and mixed-precision models.
-            quant_config: QuantizationConfig of how the model should be quantized.
-            representative_data_gen: Dataset to use for retrieving images for the models inputs.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
-            disable_activation_for_metric: Whether to disable activation quantization when computing the MP metric.
-            hessian_info_service: HessianInfoService to fetch information based on Hessian-approximation.
-
-        Returns:
-            A function that computes the metric.
-        """
-
-        raise NotImplementedError(f'{self.__class__.__name__} has to implement the '
-                             f'framework\'s get_sensitivity_evaluator method.')  # pragma: no cover
-
     def get_node_prior_info(self, node: BaseNode,
                             fw_info: FrameworkInfo,
                             graph: Graph) -> NodePriorInfo:

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -25,6 +25,7 @@
     MixedPrecisionSearchManager
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import \
     ResourceUtilization
+from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
 from model_compression_toolkit.core.common.mixed_precision.solution_refinement_procedure import \
     greedy_solution_refinement_procedure
 
@@ -78,11 +79,12 @@ def search_bit_width(graph: Graph,
 
     # Set Sensitivity Evaluator for MP search. It should always work with the original MP graph,
     # even if a virtual graph was created (and is used only for BOPS utilization computation purposes)
-    se = fw_impl.get_sensitivity_evaluator(
+    se = SensitivityEvaluation(
         graph,
         mp_config,
         representative_data_gen=representative_data_gen,
         fw_info=fw_info,
+        fw_impl=fw_impl,
         disable_activation_for_metric=disable_activation_for_metric,
         hessian_info_service=hessian_info_service)
 

model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py

@@ -20,6 +20,8 @@
 from model_compression_toolkit.constants import AXIS
 from model_compression_toolkit.core import FrameworkInfo, MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core.common import Graph, BaseNode
+from model_compression_toolkit.core.common.mixed_precision.set_layer_to_bitwidth import \
+    set_activation_quant_layer_to_bitwidth, set_weights_quant_layer_to_bitwidth
 from model_compression_toolkit.core.common.quantization.node_quantization_config import ActivationQuantizationMode
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.common.similarity_analyzer import compute_kl_divergence
@@ -41,7 +43,6 @@ def __init__(self,
                  representative_data_gen: Callable,
                  fw_info: FrameworkInfo,
                  fw_impl: Any,
-                 set_layer_to_bitwidth: Callable,
                  disable_activation_for_metric: bool = False,
                  hessian_info_service: HessianInfoService = None
                  ):
@@ -63,8 +64,6 @@ def __init__(self,
             quant_config: MP Quantization configuration for how the graph should be quantized.
             representative_data_gen: Dataset used for getting batches for inference.
             fw_impl: FrameworkImplementation object with a specific framework methods implementation.
-            set_layer_to_bitwidth: A fw-dependent function that allows to configure a configurable MP model
-                    with a specific bit-width configuration.
             disable_activation_for_metric: Whether to disable activation quantization when computing the MP metric.
             hessian_info_service: HessianInfoService to fetch Hessian approximation information.
 
@@ -74,7 +73,6 @@ def __init__(self,
         self.representative_data_gen = representative_data_gen
         self.fw_info = fw_info
         self.fw_impl = fw_impl
-        self.set_layer_to_bitwidth = set_layer_to_bitwidth
         self.disable_activation_for_metric = disable_activation_for_metric
         if self.quant_config.use_hessian_based_scores:
             if not isinstance(hessian_info_service, HessianInfoService):
@@ -307,7 +305,13 @@ def _configure_node_bitwidth(self,
                 f"Matching layers for node {node_name} not found in the mixed precision model configuration.")  # pragma: no cover
 
         for current_layer in layers_to_config:
-            self.set_layer_to_bitwidth(current_layer, mp_model_configuration[node_idx_to_configure])
+            if isinstance(current_layer, self.fw_impl.activation_quant_layer_cls):
+                set_activation_quant_layer_to_bitwidth(current_layer, mp_model_configuration[node_idx_to_configure],
+                                                       self.fw_impl)
+            else:
+                assert isinstance(current_layer, self.fw_impl.weights_quant_layer_cls)
+                set_weights_quant_layer_to_bitwidth(current_layer, mp_model_configuration[node_idx_to_configure],
+                                                    self.fw_impl)
 
     def _compute_points_distance(self,
                                  baseline_tensors: List[Any],

model_compression_toolkit/core/common/mixed_precision/set_layer_to_bitwidth.py

@@ -12,39 +12,46 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Any
+from typing import Any, Optional, TYPE_CHECKING
 
+if TYPE_CHECKING:
+    from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 
-def set_layer_to_bitwidth(quantization_layer: Any,
-                          bitwidth_idx: int,
-                          weights_quantizer_type: type,
-                          activation_quantizer_type: type,
-                          weights_quant_layer_type: type,
-                          activation_quant_layer_type: type):
+
+def set_activation_quant_layer_to_bitwidth(quantization_layer: Any,
+                                           bitwidth_idx: Optional[int],
+                                           fw_impl: 'FrameworkImplementation'):
     """
-    Configures a layer's configurable quantizer to work with a different bit-width.
+    Configures a layer's configurable activation quantizer to work with a different bit-width.
     The bit-width_idx is the index of the actual quantizer the quantizer object in the quantization_layer wraps/holds.
 
     Args:
         quantization_layer: Layer to change its bit-width.
-        bitwidth_idx: Index of the bit-width the layer should work with.
-        weights_quantizer_type: A class of weights quantizer with configurable bitwidth options.
-        activation_quantizer_type: A class of activation quantizer with configurable bitwidth options.
-        weights_quant_layer_type: A class of a weights layer wrapper.
-        activation_quant_layer_type: A class of an activation quantization holder.
+        bitwidth_idx: Index of the bit-width the layer should work with, or None to disable quantization.
+        fw_impl: framework implementation object.
     """
+    assert isinstance(quantization_layer, fw_impl.activation_quant_layer_cls)
+    # TODO irena enable after float mp
+    # assert isinstance(quantization_layer.activation_holder_quantizer, fw_impl.configurable_activation_quantizer_cls)
+    quantization_layer.activation_holder_quantizer.set_active_activation_quantizer(bitwidth_idx)
+
 
-    if isinstance(quantization_layer, weights_quant_layer_type):
-        for _, quantizer in quantization_layer.weights_quantizers.items():
-            if isinstance(quantizer, weights_quantizer_type):
-                # Setting bitwidth only for configurable layers. There might be wrapped layers that aren't configurable,
-                # for instance, if only activations are quantized with mixed precision and weights are quantized with
-                # fixed precision
-                quantizer.set_weights_bit_width_index(bitwidth_idx)
+def set_weights_quant_layer_to_bitwidth(quantization_layer: Any,
+                                        bitwidth_idx: Optional[int],
+                                        fw_impl: 'FrameworkImplementation'):
+    """
+    Configures a layer's configurable weights quantizer to work with a different bit-width.
+    The bit-width_idx is the index of the actual quantizer the quantizer object in the quantization_layer wraps/holds.
 
-    if isinstance(quantization_layer, activation_quant_layer_type):
-        if isinstance(quantization_layer.activation_holder_quantizer, activation_quantizer_type):
-            # Setting bitwidth only for configurable layers. There might be activation layers that isn't configurable,
-            # for instance, if only weights are quantized with mixed precision and activation are quantized with
-            # fixed precision
-            quantization_layer.activation_holder_quantizer.set_active_activation_quantizer(bitwidth_idx)
+    Args:
+        quantization_layer: Layer to change its bit-width.
+        bitwidth_idx: Index of the bit-width the layer should work with, or None to disable quantization.
+        fw_impl: framework implementation object.
+    """
+    assert isinstance(quantization_layer, fw_impl.weights_quant_layer_cls)
+    configurable_quantizers = [q for q in quantization_layer.weights_quantizers.values()
+                               if isinstance(q, fw_impl.configurable_weights_quantizer_cls)]
+    # TODO irena enable after float mp
+    # assert configurable_quantizers
+    for quantizer in configurable_quantizers:
+        quantizer.set_weights_bit_width_index(bitwidth_idx)

model_compression_toolkit/core/keras/keras_implementation.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 from functools import partial
-from typing import List, Any, Tuple, Callable, Dict, Union, Generator
+from typing import List, Any, Tuple, Callable, Union, Generator
 
 import numpy as np
 import tensorflow as tf
@@ -22,7 +22,7 @@
 
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
-from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianMode, HessianInfoService
+from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianMode
 from model_compression_toolkit.core.keras.data_util import data_gen_to_dataloader
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.remove_identity import RemoveIdentity
 from model_compression_toolkit.core.keras.hessian.activation_hessian_scores_calculator_keras import \
@@ -35,8 +35,6 @@
 from model_compression_toolkit.exporter.model_wrapper.keras.builder.node_to_quantizer import \
     get_weights_quantizer_for_node, get_activations_quantizer_for_node
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
-from model_compression_toolkit.core.common.mixed_precision.set_layer_to_bitwidth import set_layer_to_bitwidth
 from model_compression_toolkit.core.common.similarity_analyzer import compute_kl_divergence, compute_cs, compute_mse
 from model_compression_toolkit.core.keras.constants import ACTIVATION, SOFTMAX, SIGMOID, ARGMAX, LAYER_NAME, \
     COMBINED_NMS
@@ -61,7 +59,7 @@
     from keras.layers import Dense, Activation, Conv2D, DepthwiseConv2D, Conv2DTranspose, Concatenate, Add   # pragma: no cover
     from keras.layers.core import TFOpLambda   # pragma: no cover
 
-from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfig
+from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
@@ -95,7 +93,7 @@
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.scale_equalization import \
     ScaleEqualization, ScaleEqualizationWithPad, ScaleEqualizationMidActivation, ScaleEqualizationMidActivationWithPad
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.separableconv_decomposition import \
-    SeparableConvDecomposition, DEPTH_MULTIPLIER
+    SeparableConvDecomposition
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.shift_negative_activation import \
     keras_apply_shift_negative_correction
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.dwconv_to_conv import DwconvToConv
@@ -110,9 +108,10 @@ class KerasImplementation(FrameworkImplementation):
     """
     A class with implemented methods to support optimizing Keras models.
     """
-
-    def __init__(self):
-        super().__init__()
+    weights_quant_layer_cls = KerasQuantizationWrapper
+    activation_quant_layer_cls = KerasActivationQuantizationHolder
+    configurable_weights_quantizer_cls = ConfigurableWeightsQuantizer
+    configurable_activation_quantizer_cls = ConfigurableActivationQuantizer
 
     @property
     def constants(self):
@@ -401,42 +400,6 @@ def get_substitutions_after_second_moment_correction(self, quant_config: Quantiz
             substitutions_list.append(keras_batchnorm_refusing())
         return substitutions_list
 
-    def get_sensitivity_evaluator(self,
-                                  graph: Graph,
-                                  quant_config: MixedPrecisionQuantizationConfig,
-                                  representative_data_gen: Callable,
-                                  fw_info: FrameworkInfo,
-                                  disable_activation_for_metric: bool = False,
-                                  hessian_info_service: HessianInfoService = None) -> SensitivityEvaluation:
-        """
-        Creates and returns an object which handles the computation of a sensitivity metric for a mixed-precision
-        configuration (comparing to the float model).
-
-        Args:
-            graph: Graph to build its float and mixed-precision models.
-            quant_config: QuantizationConfig of how the model should be quantized.
-            representative_data_gen: Dataset to use for retrieving images for the models inputs.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
-            disable_activation_for_metric: Whether to disable activation quantization when computing the MP metric.
-            hessian_info_service: HessianScoresService to fetch scores based on a Hessian-approximation for the float model.
-
-        Returns:
-            A SensitivityEvaluation object.
-        """
-
-        return SensitivityEvaluation(graph=graph,
-                                     quant_config=quant_config,
-                                     representative_data_gen=representative_data_gen,
-                                     fw_info=fw_info,
-                                     fw_impl=self,
-                                     set_layer_to_bitwidth=partial(set_layer_to_bitwidth,
-                                                                   weights_quantizer_type=ConfigurableWeightsQuantizer,
-                                                                   activation_quantizer_type=ConfigurableActivationQuantizer,
-                                                                   weights_quant_layer_type=KerasQuantizationWrapper,
-                                                                   activation_quant_layer_type=KerasActivationQuantizationHolder),
-                                     disable_activation_for_metric=disable_activation_for_metric,
-                                     hessian_info_service=hessian_info_service)
-
     def get_node_prior_info(self,
                             node: BaseNode,
                             fw_info: FrameworkInfo,