remove weights_quantization_fn from WeightsAttrQuantizationConfig

Author: irenab

model_compression_toolkit/core/common/mixed_precision/configurable_quantizer_utils.py

@@ -16,10 +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
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import (get_activation_quantizer,
+                                                                                          get_weights_quantization_fn)
 
 
 def verify_candidates_descending_order(node_q_cfg: List[CandidateNodeQuantizationConfig],
@@ -79,13 +79,13 @@ def init_quantized_weights(node_q_cfg: List[CandidateNodeQuantizationConfig],
     quantized_weights = []
     for qc in node_q_cfg:
         qc_weights_attr = qc.weights_quantization_cfg.get_attr_config(kernel_attr)
-        q_weight = qc_weights_attr.weights_quantization_fn(float_weights,
-                                                           qc_weights_attr.weights_n_bits,
-                                                           True,
-                                                           qc_weights_attr.weights_quantization_params,
-                                                           qc_weights_attr.weights_per_channel_threshold,
-                                                           qc_weights_attr.weights_channels_axis[
-                                                               0])  # output channel axis
+        weights_quantization_fn = get_weights_quantization_fn(qc_weights_attr.weights_quantization_method)
+        q_weight = weights_quantization_fn(float_weights,
+                                           qc_weights_attr.weights_n_bits,
+                                           True,
+                                           qc_weights_attr.weights_quantization_params,
+                                           qc_weights_attr.weights_per_channel_threshold,
+                                           qc_weights_attr.weights_channels_axis[0])  # output channel axis
 
         quantized_weights.append(fw_tensor_convert_func(q_weight))
 

model_compression_toolkit/core/common/quantization/node_quantization_config.py

@@ -226,8 +226,6 @@ def __init__(self,
             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
         self.weights_quantization_method = weights_attr_cfg.weights_quantization_method
@@ -316,8 +314,7 @@ def __eq__(self, other: Any) -> bool:
         if not isinstance(other, WeightsAttrQuantizationConfig):
             return False  # pragma: no cover
 
-        return self.weights_quantization_fn == other.weights_quantization_fn and \
-               self.weights_quantization_params_fn == other.weights_quantization_params_fn and \
+        return self.weights_quantization_params_fn == other.weights_quantization_params_fn and \
                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 \
@@ -327,8 +324,7 @@ def __eq__(self, other: Any) -> bool:
                self.l_p_value == other.l_p_value
 
     def __hash__(self):
-        return hash((self.weights_quantization_fn,
-                     self.weights_quantization_params_fn,
+        return hash((self.weights_quantization_params_fn,
                      self.weights_channels_axis,
                      self.weights_error_method,
                      self.weights_quantization_method,

model_compression_toolkit/core/common/quantization/quantize_node.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
-
+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.graph.base_node import BaseNode
 from model_compression_toolkit.core.common.quantization.node_quantization_config import WeightsAttrQuantizationConfig
@@ -47,11 +46,12 @@ def get_quantized_weights_attr_by_qc(attr_name: str,
         output_channels_axis = None
 
     Logger.debug(f'quantizing layer {n.name} attribute {attr_name} with {weights_qc.weights_n_bits} bits')
-    quantized_kernel = weights_qc.weights_quantization_fn(n.get_weights_by_keys(attr_name),
-                                                          n_bits=weights_qc.weights_n_bits,
-                                                          signed=True,
-                                                          quantization_params=weights_qc.weights_quantization_params,
-                                                          per_channel=weights_qc.weights_per_channel_threshold,
-                                                          output_channels_axis=output_channels_axis)
+    weights_quantization_fn = get_weights_quantization_fn(weights_qc.weights_quantization_method)
+    quantized_kernel = weights_quantization_fn(n.get_weights_by_keys(attr_name),
+                                               n_bits=weights_qc.weights_n_bits,
+                                               signed=True,
+                                               quantization_params=weights_qc.weights_quantization_params,
+                                               per_channel=weights_qc.weights_per_channel_threshold,
+                                               output_channels_axis=output_channels_axis)
 
     return quantized_kernel, channels_axis

tests_pytest/keras_tests/unit_tests/core/mixed_precision/test_set_quant_layer_to_bitwidth.py

@@ -49,8 +49,12 @@ def quant_factory(nbits, *args, **kwargs):
             assert np.allclose(x*abits[ind], y)
 
     @pytest.mark.parametrize('ind', [None, 0, 1, 2])
-    def test_configure_weights(self, ind):
+    def test_configure_weights(self, ind, mocker):
         """ Test correct weights quantizer is set and applied. """
+        def quant_factory(*args, **kwargs):
+            return lambda x, nbits, *args: x * nbits
+        mocker.patch('model_compression_toolkit.core.common.mixed_precision.configurable_quantizer_utils.'
+                     'get_weights_quantization_fn', quant_factory)
         inp = keras.layers.Input(shape=(16, 16, 3))
         out = keras.layers.Conv2D(8, kernel_size=5)(inp)
         model = keras.Model(inp, out)
@@ -63,7 +67,6 @@ def test_configure_weights(self, ind):
         for qc in qcs:
             attr_cfg = qc.weights_quantization_cfg.get_attr_config(KERNEL)
             attr_cfg.weights_channels_axis = (0,)
-            attr_cfg.weights_quantization_fn = lambda x, nbits, *args: x*nbits
         quantizer = ConfigurableWeightsQuantizer(
             node_q_cfg=qcs,
             float_weights=inner_layer.kernel.numpy(),

tests_pytest/pytorch_tests/unit_tests/core/mixed_precision/test_set_quant_layer_to_bitwidth.py

@@ -52,8 +52,12 @@ def quant_factory(nbits, *args, **kwargs):
             assert torch.allclose(x*abits[ind], y)
 
     @pytest.mark.parametrize('ind', [None, 0, 1, 2])
-    def test_configure_weights(self, ind):
+    def test_configure_weights(self, ind, mocker):
         """ Test correct weights quantizer is set and applied. """
+        def quant_factory(*args, **kwargs):
+            return lambda x, nbits, *args: x * nbits
+        mocker.patch('model_compression_toolkit.core.common.mixed_precision.configurable_quantizer_utils.'
+                     'get_weights_quantization_fn', quant_factory)
         inner_layer = torch.nn.Conv2d(3, 8, kernel_size=5).to(get_working_device())
         orig_weight = inner_layer.weight.clone()
         orig_bias = inner_layer.bias.clone()
@@ -63,7 +67,6 @@ def test_configure_weights(self, ind):
         for qc in qcs:
             attr_cfg = qc.weights_quantization_cfg.get_attr_config(KERNEL)
             attr_cfg.weights_channels_axis = (0,)
-            attr_cfg.weights_quantization_fn = lambda x, nbits, *args: x*nbits
 
         quantizer = ConfigurableWeightsQuantizer(
             node_q_cfg=qcs,