separate configuring activation and weight quantizers

Author: irenab

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/set_layer_to_bitwidth.py

@@ -12,39 +12,45 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Any
+import typing
+from typing import Any, Optional
 
+if typing.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)
+    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)]
+    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,

model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -26,14 +26,12 @@
 
 import model_compression_toolkit.core.pytorch.constants as pytorch_constants
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
-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
 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.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.hessian import HessianScoresRequest, HessianMode
 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.similarity_analyzer import compute_mse, compute_kl_divergence, compute_cs
@@ -112,6 +110,10 @@ class PytorchImplementation(FrameworkImplementation):
     """
     A class with implemented methods to support optimizing Pytorch models.
     """
+    weights_quant_layer_cls = PytorchQuantizationWrapper,
+    activation_quant_layer_cls = PytorchActivationQuantizationHolder
+    configurable_weights_quantizer_cls = ConfigurableWeightsQuantizer
+    configurable_activation_quantizer_cls = ConfigurableActivationQuantizer
 
     def __init__(self):
         super().__init__()
@@ -397,43 +399,6 @@ def get_substitutions_after_second_moment_correction(self, quant_config: Quantiz
             substitutions_list.append(pytorch_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 approximations of the hessian scores 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=PytorchQuantizationWrapper,
-                                                                   activation_quant_layer_type=PytorchActivationQuantizationHolder),
-                                     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,