fixes for api change

Author: irenab

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,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import typing
-from typing import Any, Optional
+from typing import Any, Optional, TYPE_CHECKING
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 
 
@@ -32,7 +31,8 @@ def set_activation_quant_layer_to_bitwidth(quantization_layer: Any,
         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)
+    # 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)
 
 
@@ -51,6 +51,7 @@ def set_weights_quant_layer_to_bitwidth(quantization_layer: Any,
     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
+    # TODO irena enable after float mp
+    # assert configurable_quantizers
     for quantizer in configurable_quantizers:
         quantizer.set_weights_bit_width_index(bitwidth_idx)

tests/keras_tests/function_tests/test_sensitivity_eval_non_suppoerted_output.py

@@ -18,6 +18,7 @@
 
 
 from model_compression_toolkit.core import MixedPrecisionQuantizationConfig
+from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
@@ -93,12 +94,13 @@ def verify_test_for_model(self, model):
 
         # Reducing the default number of samples for Mixed precision Hessian approximation
         # to allow quick execution of the test
-        se = keras_impl.get_sensitivity_evaluator(graph,
-                                                  MixedPrecisionQuantizationConfig(use_hessian_based_scores=True,
-                                                                                   num_of_images=2),
-                                                  representative_dataset,
-                                                  DEFAULT_KERAS_INFO,
-                                                  hessian_info_service=hessian_info_service)
+        se = SensitivityEvaluation(graph,
+                                   MixedPrecisionQuantizationConfig(use_hessian_based_scores=True,
+                                                                    num_of_images=2),
+                                   representative_dataset,
+                                   DEFAULT_KERAS_INFO,
+                                   keras_impl,
+                                   hessian_info_service=hessian_info_service)
 
     def test_not_supported_output_argmax(self):
         model = argmax_output_model((8, 8, 3))

tests/keras_tests/function_tests/test_set_layer_to_bitwidth.py

@@ -25,6 +25,7 @@
     MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_search_facade import search_bit_width, \
     BitWidthSearchMethod
+from model_compression_toolkit.core.common.mixed_precision.sensitivity_evaluation import SensitivityEvaluation
 from model_compression_toolkit.core.common.model_collector import ModelCollector
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_computation import \
@@ -96,10 +97,11 @@ def representative_data_gen():
 
         calculate_quantization_params(graph, fw_impl=keras_impl, repr_data_gen_fn=representative_data_gen)
 
-        keras_impl.get_sensitivity_evaluator(graph,
-                                             core_config.mixed_precision_config,
-                                             representative_data_gen,
-                                             fw_info=fw_info)
+        SensitivityEvaluation(graph,
+                              core_config.mixed_precision_config,
+                              representative_data_gen,
+                              fw_info=fw_info,
+                              fw_impl=keras_impl)
 
         cfg = search_bit_width(graph=graph,
                                fw_info=DEFAULT_KERAS_INFO,