remove weights_quantization_params_fn from WeightsAttrQuantizationConfig

Author: irenab

model_compression_toolkit/core/common/network_editors/actions.py

@@ -23,8 +23,6 @@
 
 
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
-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_fn_selection import \
     get_weights_quantization_fn
 
@@ -234,7 +232,7 @@ def apply(self, node: BaseNode, graph):
 
 class ChangeFinalWeightsQuantizationMethod(BaseAction):
     """
-    Class ChangeFinalWeightsQuantizationMethod to change a node's weights/activations quantizer function.
+    Class ChangeFinalWeightsQuantizationMethod to change a node's weights/activations quantizer method.
     """
 
     def __init__(self, attr_name: str, weights_quantization_method=None):
@@ -260,21 +258,8 @@ def apply(self, node: BaseNode, graph):
         """
 
         if self.weights_quantization_method is not None and node.final_weights_quantization_cfg is not None:
-
-            weights_quantization_params_fn = get_weights_quantization_params_fn(self.weights_quantization_method)
-
-            attr_config = node.final_weights_quantization_cfg.get_attr_config(self.attr_name)
-            attr_config.override_weights_quantization_params_fn(weights_quantization_params_fn)
-
-            weights_quantization_fn = get_weights_quantization_fn(self.weights_quantization_method)
-
-            if weights_quantization_fn is None:
-                Logger.critical('Unknown weights quantization method specified.')  # pragma: no cover
-
             attr_config = node.final_weights_quantization_cfg.get_attr_config(self.attr_name)
-            attr_config.override_weights_quantization_fn(weights_quantization_fn)
-            node.final_weights_quantization_cfg.get_attr_config(self.attr_name).weights_quantization_method = \
-                self.weights_quantization_method
+            attr_config.weights_quantization_method = self.weights_quantization_method
 
 
 class ChangeCandidatesWeightsQuantizationMethod(BaseAction):
@@ -307,18 +292,7 @@ def apply(self, node: BaseNode, graph: Graph):
 
         if self.weights_quantization_method is not None:
             for qc in node.candidates_quantization_cfg:
-
-                weights_quantization_params_fn = get_weights_quantization_params_fn(self.weights_quantization_method)
-
                 attr_qc = qc.weights_quantization_cfg.get_attr_config(self.attr_name)
-                attr_qc.override_weights_quantization_params_fn(weights_quantization_params_fn)
-
-                weights_quantization_fn = get_weights_quantization_fn(self.weights_quantization_method)
-
-                if weights_quantization_fn is None:
-                    Logger.critical('Unknown weights quantization method specified.')  # pragma: no cover
-
-                attr_qc.override_weights_quantization_fn(weights_quantization_fn)
                 attr_qc.weights_quantization_method = self.weights_quantization_method
 
 

model_compression_toolkit/core/common/quantization/node_quantization_config.py

@@ -12,14 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Callable, Any, List, Dict, TYPE_CHECKING
+from typing import Any, List, Dict, TYPE_CHECKING
 from enum import Enum, auto
-import numpy as np
 
 from model_compression_toolkit.core.common.framework_info import ChannelAxisMapping
 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
@@ -209,8 +206,6 @@ 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.
-        self.weights_quantization_params_fn = get_weights_quantization_params_fn(weights_attr_cfg.weights_quantization_method)
         self.weights_channels_axis = weights_channels_axis
         self.weights_quantization_method = weights_attr_cfg.weights_quantization_method
         self.weights_n_bits = weights_attr_cfg.weights_n_bits
@@ -227,26 +222,6 @@ def set_qc(self, qc: QuantizationConfig):
         self.weights_error_method = qc.weights_error_method
         self.l_p_value = qc.l_p_value
 
-    def override_weights_quantization_fn(self, weights_quantization_fn: Callable):
-        """
-        Override weights quantization function for the node.
-
-        Args:
-            weights_quantization_fn: Function for quantazing the weights.
-
-        """
-        self.weights_quantization_fn = weights_quantization_fn
-
-    def override_weights_quantization_params_fn(self, weights_quantization_params_fn: Callable):
-        """
-        Override weights params function for the node.
-
-        Args:
-            weights_quantization_params_fn: Function for calculating the weights params.
-
-        """
-        self.weights_quantization_params_fn = weights_quantization_params_fn
-
     def set_weights_quantization_param(self,
                                        weights_params: dict):
         """
@@ -260,31 +235,6 @@ def set_weights_quantization_param(self,
         for param_name, param_value in weights_params.items():
             self.weights_quantization_params[param_name] = param_value
 
-    def calculate_and_set_weights_params(self, tensor_data: np.ndarray, min_threshold: float):
-        """
-        Args:
-            tensor_data: Tensor content as Numpy array.
-            min_threshold: A minimal threshold to set as quantization parameter.
-
-        Returns:
-            Recalculated weights quantization params from the kernel and channel axis.
-
-        """
-        assert self.enable_weights_quantization
-        assert not (self.weights_per_channel_threshold and self.weights_channels_axis is None), \
-            "Trying to calculate threshold per channel, channel axis in None."
-        if self.weights_quantization_params_fn is not None:
-            self.set_weights_quantization_param(
-                self.weights_quantization_params_fn(tensor_data,
-                                                    p=self.l_p_value,
-                                                    n_bits=self.weights_n_bits,
-                                                    per_channel=self.weights_per_channel_threshold and self.weights_channels_axis is not None,
-                                                    channel_axis=self.weights_channels_axis.output,  # output channel axis
-                                                    min_threshold=min_threshold)[0]  # Take only first output, the q-params, as axis is already chosen.
-            )
-        else:
-            self.set_weights_quantization_param({})
-
     def __eq__(self, other: Any) -> bool:
         """
         Compares the object to another object to find if they are equal.
@@ -298,8 +248,7 @@ def __eq__(self, other: Any) -> bool:
         if not isinstance(other, WeightsAttrQuantizationConfig):
             return False  # pragma: no cover
 
-        return self.weights_quantization_params_fn == other.weights_quantization_params_fn and \
-               self.weights_channels_axis == other.weights_channels_axis and \
+        return self.weights_channels_axis == other.weights_channels_axis and \
                self.weights_quantization_method == other.weights_quantization_method and \
                self.weights_n_bits == other.weights_n_bits and \
                self.weights_per_channel_threshold == other.weights_per_channel_threshold and \
@@ -308,8 +257,7 @@ def __eq__(self, other: Any) -> bool:
                self.l_p_value == other.l_p_value
 
     def __hash__(self):
-        return hash((self.weights_quantization_params_fn,
-                     self.weights_channels_axis,
+        return hash((self.weights_channels_axis,
                      self.weights_error_method,
                      self.weights_quantization_method,
                      self.weights_n_bits,

model_compression_toolkit/core/common/quantization/quantization_params_fn_selection.py

@@ -17,7 +17,7 @@
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.lut_kmeans_params import (
     lut_kmeans_tensor, lut_kmeans_histogram)
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.symmetric_selection import (
-    symmetric_no_clipping_selection_min_max, symmetric_selection_histogram)
+    symmetric_no_clipping_selection_min_max, symmetric_selection_histogram, symmetric_selection_tensor)
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.uniform_selection import (
-    uniform_no_clipping_selection_min_max, uniform_selection_histogram)
+    uniform_no_clipping_selection_min_max, uniform_selection_histogram, uniform_selection_tensor)
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.outlier_filter import z_score_filter

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

@@ -27,7 +27,7 @@
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_activations_computation \
     import get_activations_qparams
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_weights_computation import \
-    get_weights_qparams
+    compute_weights_qparams
 from model_compression_toolkit.logger import Logger
 
 
@@ -119,13 +119,12 @@ def calculate_quantization_params(graph: Graph,
                             mod_attr_cfg = copy.deepcopy(attr_cfg)
                             mod_attr_cfg.weights_error_method = QuantizationErrorMethod.MSE
 
-                    weights_params, output_channels_axis = get_weights_qparams(n.get_weights_by_keys(attr),
-                                                                               candidate_qc.weights_quantization_cfg,
-                                                                               mod_attr_cfg,
-                                                                               output_channels_axis,
-                                                                               node=n,
-                                                                               hessian_info_service=hessian_info_service,
-                                                                               num_hessian_samples=num_hessian_samples)
+                    min_threshold = candidate_qc.weights_quantization_cfg.min_threshold
+                    weights_params, output_channels_axis = compute_weights_qparams(n.get_weights_by_keys(attr),
+                                                                                   mod_attr_cfg, output_channels_axis,
+                                                                                   min_threshold=min_threshold, node=n,
+                                                                                   hessian_info_service=hessian_info_service,
+                                                                                   num_hessian_samples=num_hessian_samples)
                     attr_cfg.weights_channels_axis = ChannelAxisMapping(output_channels_axis, attr_cfg.weights_channels_axis.input)
                     attr_cfg.set_weights_quantization_param(weights_params)
 

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

@@ -12,35 +12,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Dict, Any, Tuple
+from functools import partial
+from typing import Dict, Any, Tuple, Callable, TYPE_CHECKING
 
 import numpy as np
+from mct_quantizers import QuantizationMethod
 
 from model_compression_toolkit.constants import NUM_QPARAM_HESSIAN_SAMPLES
 from model_compression_toolkit.core.common.hessian import HessianInfoService
-from model_compression_toolkit.defaultdict import DefaultDict
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig, \
-    WeightsAttrQuantizationConfig
+from model_compression_toolkit.core.common.quantization.quantization_params_generation import \
+    power_of_two_selection_tensor, lut_kmeans_tensor, symmetric_selection_tensor, uniform_selection_tensor
 from model_compression_toolkit.logger import Logger
 
+if TYPE_CHECKING:
+    from model_compression_toolkit.core.common.quantization.node_quantization_config import WeightsAttrQuantizationConfig
 
-def get_weights_qparams(weights_attr_values: np.ndarray,
-                        weights_quant_config: NodeWeightsQuantizationConfig,
-                        attr_quant_config: WeightsAttrQuantizationConfig,
-                        output_channels_axis: int,
-                        node=None,
-                        hessian_info_service: HessianInfoService = None,
-                        num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> Tuple[Dict[Any, Any], int]:
+
+def compute_weights_qparams(weights_attr_values: np.ndarray,
+                            attr_quant_config: 'WeightsAttrQuantizationConfig',
+                            output_channels_axis: int,
+                            min_threshold: float,
+                            node=None,
+                            hessian_info_service: HessianInfoService = None,
+                            num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> Tuple[Dict[Any, Any], int]:
     """
     Compute thresholds to quantize a kernel according to a NodeWeightsQuantizationConfig
     instance.
 
     Args:
         weights_attr_values: Weights attribute parameter to compute the quantization thresholds for.
-        weights_quant_config: Weights quantization configuration to define how the thresholds are computed.
         attr_quant_config: A specific weights attribute quantization configuration to get its params.
         output_channels_axis: Index of the kernel output channels dimension.
+        min_threshold: Minimal threshold to use if threshold is too small.
         node: The node for which the quantization error is computed (used only with HMSE error method).
         hessian_info_service: HessianInfoService object for retrieving Hessian-based scores (used only with HMSE error method).
         num_hessian_samples: Number of samples to approximate Hessian-based scores on (used only with HMSE error method).
@@ -49,22 +52,43 @@ def get_weights_qparams(weights_attr_values: np.ndarray,
         A dictionary with the quantization threshold of the kernel.
         Selected quantization channel axis.
     """
-    if attr_quant_config.weights_quantization_params_fn is not None:
-        weights_params, output_channels_axis = attr_quant_config.weights_quantization_params_fn(
-            weights_attr_values,
-            p=attr_quant_config.l_p_value,
-            n_bits=attr_quant_config.weights_n_bits,
-            per_channel=attr_quant_config.weights_per_channel_threshold,
-            channel_axis=output_channels_axis,
-            min_threshold=weights_quant_config.min_threshold,
-            quant_error_method=attr_quant_config.weights_error_method,
-            node=node,
-            hessian_info_service=hessian_info_service,
-            num_hessian_samples=num_hessian_samples)
-    else:  # pragma: no cover
-        Logger.error(f"Requested weights quantization parameters computation for node {node.name} without providing a "
-                     f"weights_quantization_params_fn."
-                     f"Returning an empty dictionary since no quantization parameters were computed.")
-        weights_params = {}
+    params_fn = _get_weights_quantization_params_fn(attr_quant_config.weights_quantization_method)
+    weights_params, output_channels_axis = params_fn(
+        weights_attr_values,
+        p=attr_quant_config.l_p_value,
+        n_bits=attr_quant_config.weights_n_bits,
+        per_channel=attr_quant_config.weights_per_channel_threshold,
+        channel_axis=output_channels_axis,
+        min_threshold=min_threshold,
+        quant_error_method=attr_quant_config.weights_error_method,
+        node=node,
+        hessian_info_service=hessian_info_service,
+        num_hessian_samples=num_hessian_samples)
 
     return weights_params, output_channels_axis
+
+
+_weights_quant_params_fns = {
+    QuantizationMethod.POWER_OF_TWO: power_of_two_selection_tensor,
+    QuantizationMethod.SYMMETRIC: symmetric_selection_tensor,
+    QuantizationMethod.UNIFORM: uniform_selection_tensor,
+    QuantizationMethod.LUT_POT_QUANTIZER: partial(lut_kmeans_tensor, is_symmetric=False),
+    QuantizationMethod.LUT_SYM_QUANTIZER: partial(lut_kmeans_tensor, is_symmetric=True)
+}
+
+
+def _get_weights_quantization_params_fn(weights_quantization_method: QuantizationMethod) -> Callable:
+    """
+    Generate a function for finding weights quantization parameters.
+
+    Args:
+        weights_quantization_method: Which quantization method to use for weights.
+    Returns:
+        A function to find the quantization parameters.
+
+    """
+    params_fn = _weights_quant_params_fns.get(weights_quantization_method)
+    if not params_fn:
+        Logger.critical(
+            f"No parameter function found for the specified quantization method: {weights_quantization_method}")  # pragma: no cover
+    return params_fn