fixes

Author: irenab

model_compression_toolkit/core/common/framework_info.py

@@ -52,20 +52,19 @@ class FrameworkInfo(ABC):
     no_quantization_ops:Layers that should not get quantized (e.g., Reshape, Transpose, etc.)
 
     Fields:
-        activation_quantizer_factories: A mapping from QuantizationMethod to a factory function that accepts
-          activation bitwidth and a dict of quantization params, and returns the corresponding quantization function.
         kernel_channels_mapping (Dict): Dictionary from a layer to a tuple of its kernel in/out channels indices.
         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_factory_mapping: A mapping from QuantizationMethod to a factory function that accepts activation bitwidth and a dict of quantization params, and returns the corresponding quantization function.
     """
 
-    activation_quantizer_factories: Dict[QuantizationMethod, Callable[[int, dict], Callable]]
-    kernel_channels_mapping: Dict[Any, ChannelAxisMapping]
     kernel_ops_attribute_mapping: Dict[Any, str]
+    kernel_channels_mapping: Dict[Any, ChannelAxisMapping]
     out_channel_axis_mapping: Dict[Any, int]
-    _layer_min_max_mapping: Dict[Any, tuple]
+    activation_quantizer_factory_mapping: Dict[QuantizationMethod, Callable[[int, dict], Callable]]
 
+    _layer_min_max_mapping: Dict[Any, tuple]
     _default_channel_mapping = ChannelAxisMapping(None, None)
 
     @classmethod

model_compression_toolkit/core/common/mixed_precision/configurable_quantizer_utils.py

@@ -18,7 +18,7 @@
 
 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_quantizer,
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import (get_activation_quantization_fn,
                                                                                           get_weights_quantization_fn)
 
 
@@ -107,7 +107,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
-        quantizer = get_activation_quantizer(q_activation)
+        quantizer = get_activation_quantization_fn(q_activation)
         activation_quantizers.append(quantizer)
 
     return activation_quantizers

model_compression_toolkit/core/common/quantization/quantization_fn_selection.py

@@ -25,7 +25,7 @@
     symmetric_quantizer, uniform_quantizer
 
 
-def get_activation_quantizer_factory(quantization_method: QuantizationMethod) -> Callable[[int, dict], Callable]:
+def get_activation_quantization_fn_factory(quantization_method: QuantizationMethod) -> Callable[[int, dict], Callable]:
     """
     Get factory for activation quantizer.
 
@@ -35,10 +35,10 @@ def get_activation_quantizer_factory(quantization_method: QuantizationMethod) ->
     Returns:
         Factory that accepts activation bitwidth and a dict of quantization params, and returns the quantizer.
     """
-    return get_fw_info().activation_quantizer_mapping[quantization_method]
+    return get_fw_info().activation_quantizer_factory_mapping[quantization_method]
 
 
-def get_activation_quantizer(activation_quantization_cfg: NodeActivationQuantizationConfig) -> Callable:
+def get_activation_quantization_fn(activation_quantization_cfg: NodeActivationQuantizationConfig) -> Callable:
     """
     Get activation quantizer based on activation quantization configuration.
 
@@ -48,7 +48,8 @@ def get_activation_quantizer(activation_quantization_cfg: NodeActivationQuantiza
     Returns:
         Activation quantizer that accepts a tensor and returns a quantized tensor.
     """
-    quantizer_factory = get_activation_quantizer_factory(activation_quantization_cfg.activation_quantization_method)
+    quantizer_factory = get_activation_quantization_fn_factory(
+        activation_quantization_cfg.activation_quantization_method)
     quantizer = quantizer_factory(activation_quantization_cfg.activation_n_bits,
                                   activation_quantization_cfg.activation_quantization_params)
     return quantizer

model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_activations_computation.py

@@ -25,6 +25,46 @@
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationErrorMethod
 
 
+def compute_activation_qparams(activation_quant_cfg: NodeActivationQuantizationConfig,
+                               node_prior_info: NodePriorInfo,
+                               out_stats_container: BaseStatsCollector) -> Dict[str, Union[np.ndarray, float, bool]]:
+    """
+    Compute the activations params for a given node in a graph according to a params function.
+
+    Args:
+        activation_quant_cfg: node's activation quantization configuration.
+        node_prior_info: Prior info collected for the node that is being quantized.
+        out_stats_container: Tensor containing output statistics of the node.
+
+    Returns:
+        The computed activation quantization params.
+    """
+    activation_quantization_params_fn = _get_activation_quantization_params_fn(
+        activation_quant_cfg.activation_quantization_method, no_clipping=node_prior_info.is_output_bounded())
+
+    # Extract and filter histogram data from the statistics container.
+    bins_values, bins_counts = _get_histogram_data(activation_quant_cfg, out_stats_container)
+
+    # Retrieve the minimum and maximum values from the statistics container.
+    min_value, max_value = out_stats_container.get_min_max_values()
+
+    # Determine if the activations should be considered signed.
+    signed = _determine_signedness(activation_quant_cfg, node_prior_info, min_value, bins_values, bins_counts)
+
+    # Compute and return the activation quantization parameters.
+    return activation_quantization_params_fn(
+        bins_values,
+        bins_counts,
+        activation_quant_cfg.l_p_value,
+        activation_quant_cfg.activation_n_bits,
+        min_value,
+        max_value,
+        min_threshold=activation_quant_cfg.min_threshold,
+        quant_error_method=activation_quant_cfg.activation_error_method,
+        is_signed=signed
+    )
+
+
 def _get_histogram_data(
     activation_quant_cfg: NodeActivationQuantizationConfig,
     out_stats_container: BaseStatsCollector
@@ -85,46 +125,6 @@ def _determine_signedness(
     return np.any(bins_values[:-1][bins_counts > 0] < 0)
 
 
-def get_activations_qparams(activation_quant_cfg: NodeActivationQuantizationConfig,
-                            node_prior_info: NodePriorInfo,
-                            out_stats_container: BaseStatsCollector) -> Dict[str, Union[np.ndarray, float, bool]]:
-    """
-    Compute the activations params for a given node in a graph according to a params function.
-
-    Args:
-        activation_quant_cfg: node's activation quantization configuration.
-        node_prior_info: Prior info collected for the node that is being quantized.
-        out_stats_container: Tensor containing output statistics of the node.
-
-    Returns:
-        The computed activation quantization params.
-    """
-    activation_quantization_params_fn = _get_activation_quantization_params_fn(
-        activation_quant_cfg.activation_quantization_method, no_clipping=node_prior_info.is_output_bounded())
-
-    # Extract and filter histogram data from the statistics container.
-    bins_values, bins_counts = _get_histogram_data(activation_quant_cfg, out_stats_container)
-
-    # Retrieve the minimum and maximum values from the statistics container.
-    min_value, max_value = out_stats_container.get_min_max_values()
-
-    # Determine if the activations should be considered signed.
-    signed = _determine_signedness(activation_quant_cfg, node_prior_info, min_value, bins_values, bins_counts)
-
-    # Compute and return the activation quantization parameters.
-    return activation_quantization_params_fn(
-        bins_values,
-        bins_counts,
-        activation_quant_cfg.l_p_value,
-        activation_quant_cfg.activation_n_bits,
-        min_value,
-        max_value,
-        min_threshold=activation_quant_cfg.min_threshold,
-        quant_error_method=activation_quant_cfg.activation_error_method,
-        is_signed=signed
-    )
-
-
 _activation_quant_params_fns = {
     QuantizationMethod.POWER_OF_TWO: qpg.power_of_two_selection_histogram,
     QuantizationMethod.SYMMETRIC: qpg.symmetric_selection_histogram,

model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_computation.py

@@ -25,7 +25,7 @@
 from model_compression_toolkit.core.common.hessian import HessianInfoService, HessianScoresRequest, HessianMode, \
     HessianScoresGranularity
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_activations_computation \
-    import get_activations_qparams
+    import compute_activation_qparams
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_weights_computation import \
     compute_weights_qparams
 from model_compression_toolkit.logger import Logger
@@ -130,9 +130,8 @@ def calculate_quantization_params(graph: Graph,
 
             if n.is_activation_quantization_enabled():
                 # If node's activations should be quantized as well, we compute its activation quantization parameters
-                activation_params = get_activations_qparams(
-                    activation_quant_cfg=candidate_qc.activation_quantization_cfg,
-                    node_prior_info=n.prior_info,
+                activation_params = compute_activation_qparams(
+                    activation_quant_cfg=candidate_qc.activation_quantization_cfg, node_prior_info=n.prior_info,
                     out_stats_container=graph.get_out_stats_collector(n))
                 # Create a NodeQuantizationConfig containing all quantization params and attach it to the node
                 candidate_qc.activation_quantization_cfg.set_activation_quantization_param(activation_params)

model_compression_toolkit/core/common/statistics_correction/apply_second_moment_correction_to_graph.py

@@ -24,7 +24,7 @@
 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_activations_computation \
-    import get_activations_qparams
+    import compute_activation_qparams
 from model_compression_toolkit.core.common.quantization.quantize_graph_weights import quantize_graph_weights
 from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
 
@@ -52,10 +52,9 @@ def _collect_and_assign_act_threshold(graph: Graph,
 
     for n in list(graph.nodes):
         if n.is_activation_quantization_enabled():
-            activation_params = get_activations_qparams(
-                activation_quant_cfg=n.final_activation_quantization_cfg,
-                node_prior_info=n.prior_info,
-                out_stats_container=graph.get_out_stats_collector(n))
+            activation_params = compute_activation_qparams(activation_quant_cfg=n.final_activation_quantization_cfg,
+                                                           node_prior_info=n.prior_info,
+                                                           out_stats_container=graph.get_out_stats_collector(n))
             n.final_activation_quantization_cfg.set_activation_quantization_param(activation_params)
 
 

model_compression_toolkit/core/common/statistics_correction/compute_activation_bias_correction_of_graph.py

@@ -19,7 +19,7 @@
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_activation_quantizer
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_activation_quantization_fn
 
 
 def get_previous_node_with_activation_quantization(linear_node: BaseNode,
@@ -106,7 +106,7 @@ def compute_activation_bias_correction(graph: Graph,
     float_centers = calculate_bin_centers(float_bins)
 
     # Quantize the bin edges and calculate the centers of the quantized bins
-    activation_quantizer = get_activation_quantizer(prev_node_act_quant_cfg)
+    activation_quantizer = get_activation_quantization_fn(prev_node_act_quant_cfg)
     quant_bins = activation_quantizer(fw_impl.to_tensor(float_bins))
     quant_bins = fw_impl.to_numpy(quant_bins)
     quant_centers = calculate_bin_centers(quant_bins)

model_compression_toolkit/core/common/substitutions/shift_negative_activation.py

@@ -19,14 +19,14 @@
 from model_compression_toolkit.core.common.quantization.node_quantization_config import WeightsAttrQuantizationConfig, \
     ActivationQuantizationMode
 from model_compression_toolkit.core.common.quantization.quantization_fn_selection import \
-    get_activation_quantizer_factory
+    get_activation_quantization_fn_factory
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.constants import THRESHOLD, SIGNED, SHIFT_NEGATIVE_NON_LINEAR_NUM_BITS
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_activations_computation \
-    import get_activations_qparams
+    import compute_activation_qparams
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.error_functions import \
     _mse_error_histogram
 from model_compression_toolkit.core.common.quantization.quantization_params_generation import z_score_filter
@@ -327,7 +327,8 @@ def shift_negative_function(graph: Graph,
                 'float32')  # Change to type float32 to support tensorflow dtypes
             for _shift_value in _q_points:
                 _hist_bins = hist_bins.astype(np.float32) + _shift_value
-                quantizer_factory = get_activation_quantizer_factory(non_linear_node_cfg_candidate.activation_quantization_method)
+                quantizer_factory = get_activation_quantization_fn_factory(
+                    non_linear_node_cfg_candidate.activation_quantization_method)
                 fw_quant_fn = quantizer_factory(non_linear_node_cfg_candidate.activation_n_bits, qparams)
                 """
                 In SNC, when better shifting values are tested for better choice,
@@ -471,11 +472,11 @@ def update(c):
                                op2d_node=op2d_node)
 
     if non_linear_node_cfg_candidate.shift_negative_threshold_recalculation:
-        activation_param = get_activations_qparams(activation_quant_cfg=non_linear_node_cfg_candidate,
-                                                   nodes_prior_info=non_linear_node.prior_info,
-                                                   out_stats_container=graph.get_out_stats_collector(non_linear_node))
+        activation_param = compute_activation_qparams(activation_quant_cfg=non_linear_node_cfg_candidate,
+                                                      node_prior_info=non_linear_node.prior_info,
+                                                      out_stats_container=graph.get_out_stats_collector(non_linear_node))
 
-        assert activation_param.get(SIGNED) == False
+        assert activation_param.get(SIGNED) is False
         for candidate_qc in non_linear_node.candidates_quantization_cfg:
             candidate_qc.activation_quantization_cfg.set_activation_quantization_param(activation_param)
 

model_compression_toolkit/core/keras/back2framework/quantized_model_builder.py

@@ -17,7 +17,7 @@
 from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
-from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_activation_quantizer
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_activation_quantization_fn
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
 from tensorflow.python.util.object_identity import Reference as TFReference
 
@@ -57,5 +57,5 @@ def _quantize_node_activations(self,
             Output of the node.
 
         """
-        activation_quantizer = get_activation_quantizer(node.final_activation_quantization_cfg)
+        activation_quantizer = get_activation_quantization_fn(node.final_activation_quantization_cfg)
         return activation_quantizer(input_tensors)

model_compression_toolkit/core/keras/default_framework_info.py

@@ -69,6 +69,16 @@ class KerasInfo(FrameworkInfo):
                                 Dense: -1,
                                 Conv2DTranspose: -1}
 
+    """
+    Mapping from a QuantizationMethod to an activation quantizer function.
+    """
+    activation_quantizer_factory_mapping = {
+        QuantizationMethod.POWER_OF_TWO: power_of_two_quantization,
+        QuantizationMethod.SYMMETRIC: symmetric_quantization,
+        QuantizationMethod.UNIFORM: uniform_quantization,
+        QuantizationMethod.LUT_POT_QUANTIZER: activation_lut_kmean_quantizer
+    }
+
     """
     Map from an activation function name to its min/max output values (if known).
     The values are used for tensor min/max values initialization.
@@ -103,14 +113,6 @@ class KerasInfo(FrameworkInfo):
                               tf.nn.softmax: (0, SOFTMAX_THRESHOLD),
                               }
 
-    """
-    Mapping from a QuantizationMethod to an activation quantizer function.
-    """
-    activation_quantizer_mapping = {QuantizationMethod.POWER_OF_TWO: power_of_two_quantization,
-                                    QuantizationMethod.SYMMETRIC: symmetric_quantization,
-                                    QuantizationMethod.UNIFORM: uniform_quantization,
-                                    QuantizationMethod.LUT_POT_QUANTIZER: activation_lut_kmean_quantizer}
-
     @classmethod
     def get_layer_min_max(cls, layer: Any, fw_attrs: Dict) -> Tuple[float, float]:
         """