remove activation_quantization_fn from NodeActivationQuantizationCfg

Author: irenab

model_compression_toolkit/core/common/framework_info.py

@@ -52,14 +52,15 @@ class FrameworkInfo(ABC):
     no_quantization_ops:Layers that should not get quantized (e.g., Reshape, Transpose, etc.)
 
     Fields:
-        activation_quantizer_mapping (Dict[QuantizationMethod, Callable]): A dictionary mapping from QuantizationMethod to a quantization function.
+        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_mapping: Dict[QuantizationMethod, Callable]
+    activation_quantizer_factories: Dict[QuantizationMethod, Callable[[int, dict], Callable]]
     kernel_channels_mapping: Dict[Any, ChannelAxisMapping]
     kernel_ops_attribute_mapping: Dict[Any, str]
     out_channel_axis_mapping: Dict[Any, int]

model_compression_toolkit/core/common/mixed_precision/configurable_quantizer_utils.py

@@ -16,8 +16,10 @@
 
 import numpy as np
 
+from model_compression_toolkit.core.common.framework_info import get_fw_info
 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
 
 
 def verify_candidates_descending_order(node_q_cfg: List[CandidateNodeQuantizationConfig],
@@ -105,6 +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
-        activation_quantizers.append(q_activation.quantize_node_output)
+        quantizer = get_activation_quantizer(q_activation)
+        activation_quantizers.append(quantizer)
 
     return activation_quantizers

model_compression_toolkit/core/common/network_editors/actions.py

@@ -250,9 +250,6 @@ def apply(self, node: BaseNode, graph):
             node.final_activation_quantization_cfg.set_activation_quantization_params_fn(
                 activation_quantization_params_fn)
 
-            activation_quantization_fn = get_fw_info().activation_quantizer_mapping.get(self.activation_quantization_method)
-
-            node.final_activation_quantization_cfg.set_activation_quantization_fn(activation_quantization_fn)
             node.final_activation_quantization_cfg.activation_quantization_method = self.activation_quantization_method
 
 
@@ -283,15 +280,7 @@ def apply(self, node: BaseNode, graph):
             for qc in node.candidates_quantization_cfg:
                 activation_quantization_params_fn = get_activation_quantization_params_fn(
                     self.activation_quantization_method)
-
                 qc.activation_quantization_cfg.set_activation_quantization_params_fn(activation_quantization_params_fn)
-                activation_quantization_fn = get_fw_info().activation_quantizer_mapping.get(
-                    self.activation_quantization_method)
-
-                if activation_quantization_fn is None:
-                    Logger.critical('Unknown activation quantization method specified.')  # pragma: no cover
-
-                qc.activation_quantization_cfg.set_activation_quantization_fn(activation_quantization_fn)
                 qc.activation_quantization_cfg.activation_quantization_method = self.activation_quantization_method
 
 

model_compression_toolkit/core/common/quantization/node_quantization_config.py

@@ -17,16 +17,14 @@
 import numpy as np
 
 from model_compression_toolkit.core.common.framework_info import ChannelAxisMapping
-from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_weights_quantization_fn
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.quantization.quantization_params_fn_selection import \
     get_weights_quantization_params_fn
 
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
 from model_compression_toolkit.target_platform_capabilities.constants import POSITIONAL_ATTR
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import \
-    AttributeQuantizationConfig, \
-    OpQuantizationConfig
+    AttributeQuantizationConfig, OpQuantizationConfig
 
 if TYPE_CHECKING:
     from model_compression_toolkit.core.common.graph.base_node import WeightAttrT
@@ -85,18 +83,14 @@ class NodeActivationQuantizationConfig(BaseNodeQuantizationConfig):
     """
     def __init__(self,
                  op_cfg: OpQuantizationConfig,
-                 activation_quantization_fn: Callable,
-                 activation_quantization_params_fn: Callable
-                 ):
+                 activation_quantization_params_fn: Callable):
         """
 
         Args:
             op_cfg: OpQuantizationConfig of the node with quantizers types to use when creating node quantization configuration.
-            activation_quantization_fn: Function to use when quantizing the node's activations.
             activation_quantization_params_fn: Function to use when computing the threshold for quantizing a node's activations.
         """
 
-        self.activation_quantization_fn = activation_quantization_fn
         self.activation_quantization_params_fn = activation_quantization_params_fn
         self.activation_quantization_method = op_cfg.activation_quantization_method
         self.activation_n_bits = op_cfg.activation_n_bits
@@ -152,36 +146,6 @@ def quantization_preserving(self):
     def fln_quantization(self):
         return self.quant_mode == ActivationQuantizationMode.FLN_QUANT
 
-    def quantize_node_output(self,
-                             tensors: Any) -> Any:
-        """
-
-        Args:
-            tensors: framework tensor/s
-
-        Returns:
-            Framework tensor/s after applying fake quantization.
-
-        """
-        fake_quant = self.activation_quantization_fn(self.activation_n_bits,
-                                                     self.activation_quantization_params)
-
-        if fake_quant is None:
-            Logger.critical(
-                "Layer is intended to be quantized, but the fake_quant function is None.")  # pragma: no cover
-
-        return fake_quant(tensors)
-
-    def set_activation_quantization_fn(self, activation_quantization_fn: Callable):
-        """
-        Sets activation quantization function for the node.
-
-        Args:
-            activation_quantization_fn: Function for quantazing the activations.
-
-        """
-        self.activation_quantization_fn = activation_quantization_fn
-
     def set_activation_quantization_params_fn(self, activation_quantization_params_fn:Callable):
         """
         Sets activation params function for the node.
@@ -218,8 +182,7 @@ def __eq__(self, other: Any) -> bool:
         if not isinstance(other, NodeActivationQuantizationConfig):
             return False  # pragma: no cover
 
-        return self.activation_quantization_fn == other.activation_quantization_fn and \
-               self.activation_quantization_params_fn == other.activation_quantization_params_fn and \
+        return self.activation_quantization_params_fn == other.activation_quantization_params_fn and \
                self.activation_error_method == other.activation_error_method and \
                self.activation_quantization_method == other.activation_quantization_method and \
                self.activation_n_bits == other.activation_n_bits and \
@@ -234,8 +197,7 @@ def __eq__(self, other: Any) -> bool:
                self.shift_negative_threshold_recalculation == other.shift_negative_threshold_recalculation
 
     def __hash__(self):
-        return hash((self.activation_quantization_fn,
-                     self.activation_quantization_params_fn,
+        return hash((self.activation_quantization_params_fn,
                      self.activation_error_method,
                      self.activation_quantization_method,
                      self.activation_n_bits,
@@ -263,6 +225,8 @@ def __init__(self,
             weights_attr_cfg: AttributeQuantizationConfig with parameters to use when creating the node's attribute quantization config.
             weights_channels_axis: Axis to quantize a node's attribute when quantizing per-channel (if not quantizing per-channel than expecting None).
         """
+        # TODO irena remove functions.
+        from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_weights_quantization_fn
         self.weights_quantization_fn = get_weights_quantization_fn(weights_attr_cfg.weights_quantization_method)
         self.weights_quantization_params_fn = get_weights_quantization_params_fn(weights_attr_cfg.weights_quantization_method)
         self.weights_channels_axis = weights_channels_axis

model_compression_toolkit/core/common/quantization/quantization_fn_selection.py

@@ -14,15 +14,46 @@
 # ==============================================================================
 
 from collections.abc import Callable
-from functools import partial
 
 from mct_quantizers import QuantizationMethod
+
+from model_compression_toolkit.core.common.framework_info import get_fw_info
+from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeActivationQuantizationConfig
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.quantization.quantizers.lut_kmeans_quantizer import lut_kmeans_quantizer
 from model_compression_toolkit.core.common.quantization.quantizers.uniform_quantizers import power_of_two_quantizer, \
     symmetric_quantizer, uniform_quantizer
 
 
+def get_activation_quantizer_factory(quantization_method: QuantizationMethod) -> Callable[[int, dict], Callable]:
+    """
+    Get factory for activation quantizer.
+
+    Args:
+        quantization_method: quantization method for activation.
+
+    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]
+
+
+def get_activation_quantizer(activation_quantization_cfg: NodeActivationQuantizationConfig) -> Callable:
+    """
+    Get activation quantizer based on activation quantization configuration.
+
+    Args:
+        activation_quantization_cfg: activation quantization configuration.
+
+    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 = quantizer_factory(activation_quantization_cfg.activation_n_bits,
+                                  activation_quantization_cfg.activation_quantization_params)
+    return quantizer
+
+
 def get_weights_quantization_fn(weights_quantization_method: QuantizationMethod) -> Callable:
     """
     Generate a function for weight quantization.

model_compression_toolkit/core/common/statistics_correction/compute_activation_bias_correction_of_graph.py

@@ -19,6 +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
 
 
 def get_previous_node_with_activation_quantization(linear_node: BaseNode,
@@ -105,7 +106,8 @@ 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
-    quant_bins = prev_node_act_quant_cfg.quantize_node_output(fw_impl.to_tensor(float_bins))
+    activation_quantizer = get_activation_quantizer(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

@@ -18,6 +18,8 @@
 
 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
 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
@@ -325,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
-                fw_quant_fn = non_linear_node_cfg_candidate.activation_quantization_fn(non_linear_node_cfg_candidate.activation_n_bits,qparams)
+                quantizer_factory = get_activation_quantizer_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,
                 the histogram (which is a numpy object) is quantized using the non-linear node activation

model_compression_toolkit/core/keras/back2framework/quantized_model_builder.py

@@ -17,6 +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.keras.back2framework.keras_model_builder import KerasModelBuilder
 from tensorflow.python.util.object_identity import Reference as TFReference
 
@@ -56,4 +57,5 @@ def _quantize_node_activations(self,
             Output of the node.
 
         """
-        return node.final_activation_quantization_cfg.quantize_node_output(input_tensors)
+        activation_quantizer = get_activation_quantizer(node.final_activation_quantization_cfg)
+        return activation_quantizer(input_tensors)

model_compression_toolkit/core/pytorch/back2framework/quantized_model_builder.py

@@ -17,9 +17,9 @@
 
 import torch
 
-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.user_info import UserInformation
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder, \
     PytorchModel
@@ -60,7 +60,8 @@ def _quantize_node_activations(self,
         if node.is_activation_quantization_enabled():
             if isinstance(input_tensors, list):
                 input_tensors = torch.cat(input_tensors, dim=0)
-            return node.final_activation_quantization_cfg.quantize_node_output(input_tensors)
+            activation_quantizer = get_activation_quantizer(node.final_activation_quantization_cfg)
+            return activation_quantizer(input_tensors)
         return input_tensors
 
 

model_compression_toolkit/quantization_preparation/load_fqc.py

@@ -209,10 +209,8 @@ def _create_candidate(weight_channel_axis: ChannelAxisMapping,
     """
 
     # TODO irena: i think we shouldn't inject methods here, it's quantization implementation, not configuration
-    activation_quantization_fn = get_fw_info().activation_quantizer_mapping[op_cfg.activation_quantization_method]
     activation_quantization_params_fn = get_activation_quantization_params_fn(op_cfg.activation_quantization_method)
     aqc = NodeActivationQuantizationConfig(op_cfg=op_cfg,
-                                           activation_quantization_fn=activation_quantization_fn,
                                            activation_quantization_params_fn=activation_quantization_params_fn)
 
     # TODO: remove this validation and warning once enabling all attributes quantization by default