mct-nightly 2.0.0.20240413.406__py3-none-any.whl → 2.0.0.20240415.5018__py3-none-any.whl

{mct_nightly-2.0.0.20240413.406.dist-info → mct_nightly-2.0.0.20240415.5018.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.0.0.20240413.406
+Version: 2.0.0.20240415.5018
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct_nightly-2.0.0.20240413.406.dist-info → mct_nightly-2.0.0.20240415.5018.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=aBKAsetxgA8LyN7yib2REk7UPRyfK2jBb-YqBpIpCbE,1573
+model_compression_toolkit/__init__.py,sha256=pU9oIc4ZlkMr0MR9kraXEOboqcZF2lShgkyhaDHxzn0,1573
 model_compression_toolkit/constants.py,sha256=f9at1H_-vb5nvdHRmAHUco4ja4_QermK6yu0N9qbRGE,3723
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -109,7 +109,7 @@ model_compression_toolkit/core/common/quantization/quantize_graph_weights.py,sha
 model_compression_toolkit/core/common/quantization/quantize_node.py,sha256=cdzGNWfT4MRogIU8ehs0tr3lVjnzAI-jeoS9b4TwVBo,2854
 model_compression_toolkit/core/common/quantization/set_node_quantization_config.py,sha256=9BEv2l0z2trDEsr40VB8tO3ToBA_b2sd_jH9uqZ5Wo8,11503
 model_compression_toolkit/core/common/quantization/quantization_params_generation/__init__.py,sha256=eCDGwsWYLU6z7qbEVb4TozMW_nd5VEP_iCJ6PcvyEPw,1486
-model_compression_toolkit/core/common/quantization/quantization_params_generation/error_functions.py,sha256=eH3nSXPFn94ATF3dZn2HxNAGVJUWotirN6o8wwDfkLg,18165
+model_compression_toolkit/core/common/quantization/quantization_params_generation/error_functions.py,sha256=TUJuSpX8pcsIPbJ6z_YGWgD_uafqlKRJcpsTIFpjMKU,19936
 model_compression_toolkit/core/common/quantization/quantization_params_generation/lut_kmeans_params.py,sha256=HSbAlDKXZMn8BtQQGL8TnlXvO2f_2oTLXAK1khraX7g,7410
 model_compression_toolkit/core/common/quantization/quantization_params_generation/outlier_filter.py,sha256=9gnfJV89jpGwAx8ImJ5E9NjCv3lDtbyulP4OtgWb62M,1772
 model_compression_toolkit/core/common/quantization/quantization_params_generation/power_of_two_selection.py,sha256=BiwDqt5CeU6CW0Qusy3LwWhFtf2J9BvSuGMsTsG6rSw,8538
@@ -480,8 +480,8 @@ model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=7bbzqJN8ZAycVDvZr_5xC-niTAR5df8f03Kooev_pfg,3047
-mct_nightly-2.0.0.20240413.406.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.0.0.20240413.406.dist-info/METADATA,sha256=KPkoZIsVNAhmDShzs6X5LUBpcV_hvmCO9elwhFuTduw,18795
-mct_nightly-2.0.0.20240413.406.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-mct_nightly-2.0.0.20240413.406.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.0.0.20240413.406.dist-info/RECORD,,
+mct_nightly-2.0.0.20240415.5018.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.0.0.20240415.5018.dist-info/METADATA,sha256=NaQHHj_S3oEuCRqkeaaOSlRIUc6HthT3C9IjzX6p7FQ,18796
+mct_nightly-2.0.0.20240415.5018.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+mct_nightly-2.0.0.20240415.5018.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.0.0.20240415.5018.dist-info/RECORD,,

model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.0.0.20240413.000406"
+__version__ = "2.0.0.20240415.005018"

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

@@ -89,8 +89,8 @@ def _lp_error_histogram(q_bins: np.ndarray,
 
 
 def _kl_error_function(x: np.ndarray,
-                       range_min: float,
-                       range_max: float,
+                       range_min: np.ndarray,
+                       range_max: np.ndarray,
                        n_bins: int = 2048,
                        n_bits: int = 8) -> np.float32:
     """
@@ -148,7 +148,8 @@ def _kl_error_function_wrapper(x: np.ndarray,
                                range_min: np.ndarray,
                                range_max: np.ndarray,
                                n_bins: int = 2048,
-                               n_bits: int = 8) -> np.ndarray:
+                               n_bits: int = 8,
+                               per_channel: int = False) -> np.ndarray:
     """
     Computes the error function between a tensor and its quantized version for each channel.
     The error is based on the KL-divergence between the distributions.
@@ -161,6 +162,7 @@ def _kl_error_function_wrapper(x: np.ndarray,
         range_max: Array specifying the maximum bound of the quantization range for each channel.
         n_bins: Number of bins for the float histogram.
         n_bits: Number of bits used for quantization.
+        per_channel: Whether quantization is done per-channel.
 
     Returns:
         An array containing the KL-divergence between the float and quantized histograms of the tensor for each channel.
@@ -168,8 +170,11 @@ def _kl_error_function_wrapper(x: np.ndarray,
     """
 
     error_list = []
-    for j in range(x.shape[0]):  # iterate all channels of the tensor.
-        error_list.append(_kl_error_function(x[j], range_min[j], range_max[j], n_bins=n_bins, n_bits=n_bits))
+    if per_channel:
+        for j in range(x.shape[0]):  # iterate all channels of the tensor.
+            error_list.append(_kl_error_function(x[j], range_min[j], range_max[j], n_bins=n_bins, n_bits=n_bits))
+    else:
+        error_list.append(_kl_error_function(x, range_min, range_max, n_bins=n_bins, n_bits=n_bits))
     return np.asarray(error_list)
 
 
@@ -177,8 +182,8 @@ def _kl_error_histogram(q_bins: np.ndarray,
                         q_count: np.ndarray,
                         bins: np.ndarray,
                         counts: np.ndarray,
-                        range_min: float,
-                        range_max: float) -> np.float32:
+                        range_min: np.ndarray,
+                        range_max: np.ndarray) -> np.float32:
     """
     Compute the error function between a histogram to its quantized version.
     The error is computed based on the KL-divergence the distributions have.
@@ -241,8 +246,8 @@ def _kl_error_histogram(q_bins: np.ndarray,
 
 
 def _get_bins_indices_from_range(bins: np.ndarray,
-                                 range_min: float,
-                                 range_max: float) -> Tuple[int, int]:
+                                 range_min: np.ndarray,
+                                 range_max: np.ndarray) -> Tuple[int, int]:
     """
     For bins and a threshold, compute the first and last bins in between the threshold
     ranges.
@@ -262,7 +267,7 @@ def _get_bins_indices_from_range(bins: np.ndarray,
     return first_bin_idx, last_bin_idx
 
 
-def _is_range_valid(bins: np.ndarray, range_min: float, range_max: float) -> bool:
+def _is_range_valid(bins: np.ndarray, range_min: np.ndarray, range_max: np.ndarray) -> bool:
     """
     Check whether there are some bins from a numpy array of bins that are in between
     a threshold range or not.
@@ -387,15 +392,36 @@ def get_threshold_selection_tensor_error_function(quantization_method: Quantizat
 
     Returns: a Callable method that calculates the error between a tensor and a quantized tensor.
     """
+    if quant_error_method == qc.QuantizationErrorMethod.KL:
+        if axis is None:
+            # per-tensor
+            if quantization_method == QuantizationMethod.UNIFORM:
+                return lambda x, y, threshold: _kl_error_function_wrapper(x, range_min=threshold[0],
+                                                                          range_max=threshold[1],
+                                                                          n_bits=n_bits,
+                                                                          per_channel=False)
+            else:
+                return lambda x, y, threshold: _kl_error_function_wrapper(x, range_min=0 if not signed else -threshold,
+                                                                          range_max=threshold,
+                                                                          n_bits=n_bits,
+                                                                          per_channel=False)
+        else:
+            # per-channel
+            if quantization_method == QuantizationMethod.UNIFORM:
+                return lambda x, y, threshold: _kl_error_function_wrapper(x, range_min=threshold[:, 0],
+                                                                          range_max=threshold[:, 1],
+                                                                          n_bits=n_bits,
+                                                                          per_channel=True)
+            else:
+                return lambda x, y, threshold: _kl_error_function_wrapper(x, range_min=0 if not signed else -threshold,
+                                                                          range_max=threshold,
+                                                                          n_bits=n_bits,
+                                                                          per_channel=True)
 
     quant_method_error_function_mapping = {
         qc.QuantizationErrorMethod.MSE: lambda x, y, threshold: compute_mse(x, y, norm=norm, axis=axis),
         qc.QuantizationErrorMethod.MAE: lambda x, y, threshold: compute_mae(x, y, norm=norm, axis=axis),
         qc.QuantizationErrorMethod.LP: lambda x, y, threshold: compute_lp_norm(x, y, p=p, norm=norm, axis=axis),
-        qc.QuantizationErrorMethod.KL:
-            lambda x, y, threshold: _kl_error_function_wrapper(x, range_min=threshold[:,0], range_max=threshold[:,1],
-                                                       n_bits=n_bits) if quantization_method == QuantizationMethod.UNIFORM
-            else _kl_error_function_wrapper(x, range_min=0 if not signed else -threshold, range_max=threshold, n_bits=n_bits)
     }
 
     return quant_method_error_function_mapping[quant_error_method]