mct-nightly 2.1.0.20240609.524__py3-none-any.whl → 2.1.0.20240611.428__py3-none-any.whl

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

@@ -27,7 +27,7 @@ from model_compression_toolkit.constants import MIN_THRESHOLD, DEFAULT_TOL, DEFA
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import quantize_tensor, \
     reshape_tensor_for_per_channel_search, uniform_quantize_tensor, get_output_shape
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import max_power_of_two, \
-    get_tensor_max
+    get_tensor_max, get_tensor_min
 
 
 def qparams_selection_tensor_search(error_function: Callable,
@@ -56,41 +56,49 @@ def qparams_selection_tensor_search(error_function: Callable,
         signed: a flag whether the tensor is signed.
 
     Returns:
-        Optimal constrained threshold to quantize the tensor.
+        Optimal constrained threshold to quantize the tensor, and best channel axis if input channel_axis was None,
+        else return the input channel axis.
 
     """
 
-    output_shape = get_output_shape(tensor_data.shape, channel_axis)
-
-    # First threshold to check is the constrained threshold based on the tensor's maximal value.
-    tensor_max = get_tensor_max(tensor_data, per_channel, channel_axis, n_bits)
-    threshold = 2 * max_power_of_two(tensor_max, min_threshold)
+    search_axes = range(len(tensor_data.shape)) if channel_axis is None and per_channel else [channel_axis]
+    total_error_list = []
+    th_list = []
+    for _axis in search_axes:
+        output_shape = get_output_shape(tensor_data.shape, _axis)
 
-    # If the threshold is computed per-channel, we rearrange the tensor such that each sub-tensor
-    # is flattened, and we iterate over each one of them when searching for the threshold.
-    if per_channel:
-        tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, channel_axis)
+        # First threshold to check is the constrained threshold based on the tensor's maximal value.
+        tensor_max = get_tensor_max(tensor_data, per_channel, _axis, n_bits)
+        threshold = 2 * max_power_of_two(tensor_max, min_threshold)
 
-    error_list = []  # init an empty error list
-    # On each iteration a new constrained threshold which equal to half of the previous tested threshold
-    # is used for quantizing the tensor and computing the error. The error is appended to an error list, which
-    # eventually used to select the threshold with the minimal error.
-    for i in range(n_iter):
+        # Rearrange the tensor such that each sub-tensor is flattened, and we iterate over each
+        # one of them when searching for the threshold.
         if per_channel:
-            threshold_hat = (threshold / (2 ** i)).reshape([-1, 1])
-            qt = quantize_tensor(tensor_data_r, threshold_hat, n_bits, signed)
-            per_channel_error = _error_function_wrapper(error_function, tensor_data_r, qt, threshold_hat)
-
-            error_list.append(per_channel_error)
-        else:  # quantize per-tensor
-            qt = quantize_tensor(tensor_data, threshold / (2 ** i), n_bits, signed)
-            error = error_function(qt, tensor_data, threshold=threshold / (2 ** i))
-            error_list.append(error)
-
-    # Take the index of the minimal error, and use it compute the threshold which yielded it.
-    i = np.argmin(np.stack(error_list, axis=-1), axis=-1)
-
-    return np.maximum(np.reshape(threshold.flatten() / np.power(2, i), output_shape), min_threshold)
+            tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, _axis)
+
+        error_list = []  # init an empty error list
+        # On each iteration a new constrained threshold which equal to half of the previous tested threshold
+        # is used for quantizing the tensor and computing the error. The error is appended to an error list, which
+        # eventually used to select the threshold with the minimal error.
+        for i in range(n_iter):
+            if per_channel:
+                threshold_hat = (threshold / (2 ** i)).reshape([-1, 1])
+                qt = quantize_tensor(tensor_data_r, threshold_hat, n_bits, signed)
+                per_channel_error = _error_function_wrapper(error_function, tensor_data_r, qt, threshold_hat)
+                error_list.append(per_channel_error)
+            else:  # quantize per-tensor
+                qt = quantize_tensor(tensor_data, threshold / (2 ** i), n_bits, signed)
+                error = error_function(qt, tensor_data, threshold=threshold / (2 ** i))
+                error_list.append(error)
+
+        # Take the index of the minimal error, and use it compute the threshold which yielded it.
+        err_mat = np.stack(error_list, axis=-1)
+        i = np.argmin(err_mat, axis=-1)
+        th_list.append(np.maximum(np.reshape(threshold.flatten() / np.power(2, i), output_shape), min_threshold))
+        total_error_list.append(err_mat.min(axis=-1).mean())
+
+    best_axis_index = np.argmin(total_error_list)
+    return th_list[best_axis_index], search_axes[best_axis_index]
 
 
 def qparams_selection_histogram_search(error_function: Callable,
@@ -390,13 +398,12 @@ def search_dynamic_range(base_range: np.ndarray, x: np.ndarray, scalers: np.ndar
 
 def qparams_symmetric_selection_tensor_search(error_function: Callable,
                                               tensor_data: np.ndarray,
-                                              tensor_max: np.ndarray,
                                               n_bits: int,
                                               per_channel: bool = False,
                                               channel_axis: int = 1,
                                               n_iter: int = SYMMETRIC_TENSOR_PER_CHANNEL_N_ITER,
                                               min_threshold=MIN_THRESHOLD,
-                                              signed: bool = True) -> Any:
+                                              signed: bool = True) -> Tuple[np.ndarray, int]:
     """
     Search for optimal threshold (per-channel or per-tensor) for symmetric quantization of a tensor,
     using the iterative optimizer method.
@@ -404,7 +411,6 @@ def qparams_symmetric_selection_tensor_search(error_function: Callable,
     Args:
         error_function: Function to compute the error between the original and quantized tensors.
         tensor_data: Numpy array with tensor's content.
-        tensor_max: The max value of the tensor.
         n_bits: Number of bits to quantize the tensor.
         per_channel: Whether the tensor should be quantized per-channel or per-tensor.
         channel_axis: Index of output channels dimension.
@@ -417,46 +423,55 @@ def qparams_symmetric_selection_tensor_search(error_function: Callable,
 
     """
 
-    output_shape = get_output_shape(tensor_data.shape, channel_axis)
+    search_axes = range(len(tensor_data.shape)) if channel_axis is None and per_channel else [channel_axis]
+    total_error_list = []
+    th_list = []
+    for _axis in search_axes:
+        tensor_max = get_tensor_max(tensor_data, per_channel, _axis, n_bits)
+        output_shape = get_output_shape(tensor_data.shape, _axis)
 
-    # If the threshold is computed per-channel, we rearrange the tensor such that each sub-tensor
-    # is flattened, and we iterate over each one of them when searching for the threshold.
-    if per_channel:
-        tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, channel_axis)
-        max_tensor = np.maximum(min_threshold, tensor_max)
-        res = qparams_symmetric_iterative_minimization(x0=max_tensor,
-                                                       x=tensor_data_r,
-                                                       loss_fn=error_function,  # gets float_tensor, fxp_tensor, threshold
-                                                       n_bits=n_bits,
-                                                       signed=signed,
-                                                       n_intervals=SYMMETRIC_TENSOR_PER_CHANNEL_N_INTERVALS,
-                                                       n_iter=SYMMETRIC_TENSOR_PER_CHANNEL_N_ITER,
-                                                       dec_freq=SYMMETRIC_TENSOR_PER_CHANNEL_DEC_FREQ,
-                                                       per_channel=True)
-        return np.reshape(np.maximum(min_threshold, res['param']), output_shape)
-    else:
-        # quantize per-tensor
-        res = qparams_symmetric_iterative_minimization(x0=get_init_threshold(min_threshold, tensor_max),
-                                                       x=tensor_data,
-                                                       loss_fn=error_function,
-                                                       n_bits=n_bits,
-                                                       signed=signed,
-                                                       n_intervals=SYMMETRIC_TENSOR_N_INTERVALS,
-                                                       n_iter=SYMMETRIC_TENSOR_N_ITER,
-                                                       dec_freq=SYMMETRIC_TENSOR_DEC_FREQ,
-                                                       per_channel=False)
-
-        return max(min_threshold, res['param'])
+        if per_channel:
+            # Rearrange the tensor such that each sub-tensor is flattened, and we iterate
+            # over each one of them when searching for the threshold.
+            tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, _axis)
+            max_tensor = np.maximum(min_threshold, tensor_max)
+            res = qparams_symmetric_iterative_minimization(x0=max_tensor,
+                                                           x=tensor_data_r,
+                                                           loss_fn=error_function,  # gets float_tensor, fxp_tensor, threshold
+                                                           n_bits=n_bits,
+                                                           signed=signed,
+                                                           n_intervals=SYMMETRIC_TENSOR_PER_CHANNEL_N_INTERVALS,
+                                                           n_iter=SYMMETRIC_TENSOR_PER_CHANNEL_N_ITER,
+                                                           dec_freq=SYMMETRIC_TENSOR_PER_CHANNEL_DEC_FREQ,
+                                                           per_channel=True)
+            th = np.reshape(np.maximum(min_threshold, res['param']), output_shape)
+        else:
+            # quantize per-tensor
+            res = qparams_symmetric_iterative_minimization(x0=get_init_threshold(min_threshold, tensor_max),
+                                                           x=tensor_data,
+                                                           loss_fn=error_function,
+                                                           n_bits=n_bits,
+                                                           signed=signed,
+                                                           n_intervals=SYMMETRIC_TENSOR_N_INTERVALS,
+                                                           n_iter=SYMMETRIC_TENSOR_N_ITER,
+                                                           dec_freq=SYMMETRIC_TENSOR_DEC_FREQ,
+                                                           per_channel=False)
+            th = max(min_threshold, res['param'])
+
+        total_error_list.append(res['loss'].mean())
+        th_list.append(th)
+
+    best_axis_index = np.argmin(total_error_list)
+    return th_list[best_axis_index], search_axes[best_axis_index]
 
 
 def qparams_uniform_selection_tensor_search(error_function: Callable,
                                             tensor_data: np.ndarray,
-                                            tensor_min: np.ndarray,
-                                            tensor_max: np.ndarray,
                                             n_bits: int,
                                             per_channel: bool = False,
                                             channel_axis: int = 1,
-                                            n_iter: int = UNIFORM_TENSOR_PER_CHANNEL_N_ITER) -> Any:
+                                            n_iter: int = UNIFORM_TENSOR_PER_CHANNEL_N_ITER,
+                                            ) -> Tuple[Tuple[np.ndarray, np.ndarray], int]:
     """
     Search for optimal quantization range (per-channel or per-tensor) for uniform quantization of a tensor,
     using the iterative optimizer method and built-in scale factors
@@ -465,8 +480,6 @@ def qparams_uniform_selection_tensor_search(error_function: Callable,
     Args:
         error_function: Function to compute the error between the original and quantized tensors.
         tensor_data: Numpy array with tensor's content.
-        tensor_min: The min value of the tensor.
-        tensor_max: The max value of the tensor.
         n_bits: Number of bits to quantize the tensor.
         per_channel: Whether the tensor should be quantized per-channel or per-tensor.
         channel_axis: Index of output channels dimension.
@@ -477,17 +490,22 @@ def qparams_uniform_selection_tensor_search(error_function: Callable,
 
     """
 
-    output_shape = get_output_shape(tensor_data.shape, channel_axis)
+    search_axes = range(len(tensor_data.shape)) if channel_axis is None and per_channel else [channel_axis]
+    total_error_list = []
+    th_list = []
+    for _axis in search_axes:
+        tensor_min = get_tensor_min(tensor_data, per_channel, _axis)
+        tensor_max = get_tensor_max(tensor_data, per_channel, _axis, n_bits, is_uniform_quantization=True)
+        output_shape = get_output_shape(tensor_data.shape, _axis)
 
-    alpha = np.linspace(BOTTOM_FACTOR, UPPER_FACTOR, UNIFORM_TENSOR_N_SAMPLES)
-    beta = np.linspace(BOTTOM_FACTOR, UPPER_FACTOR, UNIFORM_TENSOR_N_SAMPLES)
-    scalers = np.asarray(list(itertools.product(alpha, beta)))
+        alpha = np.linspace(BOTTOM_FACTOR, UPPER_FACTOR, UNIFORM_TENSOR_N_SAMPLES)
+        beta = np.linspace(BOTTOM_FACTOR, UPPER_FACTOR, UNIFORM_TENSOR_N_SAMPLES)
+        scalers = np.asarray(list(itertools.product(alpha, beta)))
 
-    # If the threshold is computed per-channel, we rearrange the tensor such that each sub-tensor
-    # is flattened, and we iterate over each one of them when searching for the threshold.
-    if per_channel:
+        # Rearrange the tensor such that each sub-tensor is flattened, and we iterate over
+        # each one of them when searching for the threshold.
         if per_channel:
-            tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, channel_axis)
+            tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, _axis)
             tensor_min_max = np.column_stack([tensor_min.flatten(), tensor_max.flatten()])
             res = iterative_uniform_dynamic_range_search(x0=tensor_min_max,
                                                          x=tensor_data_r,
@@ -496,18 +514,21 @@ def qparams_uniform_selection_tensor_search(error_function: Callable,
                                                          n_bits=n_bits,
                                                          n_iter=UNIFORM_TENSOR_PER_CHANNEL_N_ITER,
                                                          per_channel=True)
-            return np.reshape(res['param'][:, 0], output_shape), np.reshape(res['param'][:, 1], output_shape)
-    else:
-        # quantize per-tensor
-        pass
-        res = iterative_uniform_dynamic_range_search(x0=np.array([tensor_min, tensor_max]),
-                                                     x=tensor_data,
-                                                     scalers=scalers,
-                                                     loss_fn=error_function,
-                                                     n_bits=n_bits,
-                                                     n_iter=UNIFORM_TENSOR_N_ITER,
-                                                     per_channel=False)
-        return res['param']
+            th_list.append((np.reshape(res['param'][:, 0], output_shape), np.reshape(res['param'][:, 1], output_shape)))
+        else:
+            # quantize per-tensor
+            res = iterative_uniform_dynamic_range_search(x0=np.array([tensor_min, tensor_max]),
+                                                         x=tensor_data,
+                                                         scalers=scalers,
+                                                         loss_fn=error_function,
+                                                         n_bits=n_bits,
+                                                         n_iter=UNIFORM_TENSOR_N_ITER,
+                                                         per_channel=False)
+            th_list.append(tuple(np.split(res['param'], 2)))
+        total_error_list.append(res['loss'].mean())
+
+    best_axis_index = np.argmin(total_error_list)
+    return th_list[best_axis_index], search_axes[best_axis_index]
 
 
 def qparams_symmetric_selection_histogram_search(error_function: Callable,

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Dict, Any
+from typing import Dict, Any, Tuple
 
 import numpy as np
 
@@ -34,7 +34,7 @@ def get_weights_qparams(weights_attr_values: np.ndarray,
                         output_channels_axis: int,
                         node=None,
                         hessian_info_service: HessianInfoService = None,
-                        num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> Dict[Any, Any]:
+                        num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> Tuple[Dict[Any, Any], int]:
     """
     Compute thresholds to quantize a kernel according to a NodeWeightsQuantizationConfig
     instance.
@@ -50,22 +50,24 @@ def get_weights_qparams(weights_attr_values: np.ndarray,
 
     Returns:
         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 = 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 and output_channels_axis is not None,
-                                                                          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)
+        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:
         weights_params = {}
 
-    return weights_params
+    return weights_params, output_channels_axis
 
 
 def _get_kernel_channels_mapping(fw_info:FrameworkInfo,

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

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import numpy as np
+from typing import Union, Tuple, Dict
 
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
 from model_compression_toolkit.constants import MIN_THRESHOLD, THRESHOLD, NUM_QPARAM_HESSIAN_SAMPLES
@@ -25,6 +26,8 @@ from model_compression_toolkit.core.common.quantization.quantization_params_gene
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import \
     get_tensor_max
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+from model_compression_toolkit.core.common.similarity_analyzer import compute_mse
+from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import quantize_tensor
 
 
 def symmetric_selection_tensor(tensor_data: np.ndarray,
@@ -37,7 +40,8 @@ def symmetric_selection_tensor(tensor_data: np.ndarray,
                                quant_error_method: qc.QuantizationErrorMethod = qc.QuantizationErrorMethod.MSE,
                                node=None,
                                hessian_info_service: HessianInfoService = None,
-                               num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> dict:
+                               num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES,
+                               ) -> Tuple[Dict[str, np.ndarray], int]:
     """
     Compute the optimal threshold based on the provided QuantizationErrorMethod to quantize the tensor.
     Different search is applied, depends on the value of the selected QuantizationErrorMethod.
@@ -47,7 +51,7 @@ def symmetric_selection_tensor(tensor_data: np.ndarray,
         p: p-norm to use for the Lp-norm distance.
         n_bits: Number of bits to quantize the tensor.
         per_channel: Whether the quantization should be per-channel or not.
-        channel_axis: Output channel index.
+        channel_axis: Output channel index. if None, search for best axis.
         n_iter: Number of iterations to search for the optimal threshold (not used for this method).
         min_threshold: Minimal threshold to use if threshold is too small (not used for this method).
         quant_error_method: an error function to optimize the parameters' selection accordingly.
@@ -57,12 +61,24 @@ def symmetric_selection_tensor(tensor_data: np.ndarray,
 
     Returns:
         Optimal threshold to quantize the tensor in a symmetric manner.
+        Selected quantization channel axis.
     """
 
-    tensor_max = get_tensor_max(tensor_data, per_channel, channel_axis, n_bits)
-
     if quant_error_method == qc.QuantizationErrorMethod.NOCLIPPING:
-        threshold = get_init_threshold(min_threshold, tensor_max, per_channel)
+        if channel_axis is None and per_channel:
+            total_error_list = []
+            th_list = []
+            for _axis in range(len(tensor_data.shape)):
+                tensor_max = get_tensor_max(tensor_data, per_channel, _axis, n_bits)
+                threshold = get_init_threshold(min_threshold, tensor_max, per_channel)
+                q_tensor_data = quantize_tensor(tensor_data, threshold, n_bits, True)
+                total_error_list.append(compute_mse(tensor_data, q_tensor_data, norm=True))
+                th_list.append(threshold)
+            channel_axis = np.argmin(total_error_list)
+            threshold = th_list[channel_axis]
+        else:
+            tensor_max = get_tensor_max(tensor_data, per_channel, channel_axis, n_bits)
+            threshold = get_init_threshold(min_threshold, tensor_max, per_channel)
     else:
         signed = True  # weights are always signed
         axis = -1 if per_channel else None
@@ -71,15 +87,14 @@ def symmetric_selection_tensor(tensor_data: np.ndarray,
                                                                        signed=signed, node=node,
                                                                        hessian_info_service=hessian_info_service,
                                                                        num_hessian_samples=num_hessian_samples)
-        threshold = qparams_symmetric_selection_tensor_search(error_function,
-                                                              tensor_data,
-                                                              tensor_max,
-                                                              n_bits,
-                                                              per_channel,
-                                                              channel_axis,
-                                                              min_threshold=min_threshold,
-                                                              signed=signed)
-    return {THRESHOLD: threshold}
+        threshold, channel_axis = qparams_symmetric_selection_tensor_search(error_function,
+                                                                            tensor_data,
+                                                                            n_bits,
+                                                                            per_channel,
+                                                                            channel_axis,
+                                                                            min_threshold=min_threshold,
+                                                                            signed=signed)
+    return {THRESHOLD: threshold}, channel_axis
 
 
 def symmetric_selection_histogram(bins: np.ndarray,

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

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import numpy as np
+from typing import Union, Tuple, Dict
 
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
 from model_compression_toolkit.constants import MIN_THRESHOLD, RANGE_MIN, RANGE_MAX, NUM_QPARAM_HESSIAN_SAMPLES
@@ -24,6 +25,9 @@ from model_compression_toolkit.core.common.quantization.quantization_params_gene
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import get_tensor_max, \
     get_tensor_min
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+from model_compression_toolkit.core.common.similarity_analyzer import compute_mse
+from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import uniform_quantize_tensor
+
 
 def uniform_selection_tensor(tensor_data: np.ndarray,
                              p: int,
@@ -35,7 +39,8 @@ def uniform_selection_tensor(tensor_data: np.ndarray,
                              quant_error_method: qc.QuantizationErrorMethod = qc.QuantizationErrorMethod.MSE,
                              node=None,
                              hessian_info_service: HessianInfoService = None,
-                             num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> dict:
+                             num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES,
+                             ) -> Tuple[Dict[str, np.ndarray], int]:
     """
     Compute the optimal quantization range based on the provided QuantizationErrorMethod
     to uniformly quantize the tensor.
@@ -46,7 +51,7 @@ def uniform_selection_tensor(tensor_data: np.ndarray,
         p: p-norm to use for the Lp-norm distance.
         n_bits: Number of bits to quantize the tensor.
         per_channel: Whether the quantization should be per-channel or not.
-        channel_axis: Output channel index.
+        channel_axis: Output channel index. if None, search for best axis.
         n_iter: Number of iterations to search for the optimal threshold (not used for this method).
         min_threshold: Minimal threshold to use if threshold is too small (not used for this method).
         quant_error_method: an error function to optimize the range parameters' selection accordingly.
@@ -56,27 +61,48 @@ def uniform_selection_tensor(tensor_data: np.ndarray,
 
     Returns:
         Optimal quantization range to quantize the tensor uniformly.
+        Selected quantization channel axis.
     """
-    tensor_min = get_tensor_min(tensor_data, per_channel, channel_axis)
-    tensor_max = get_tensor_max(tensor_data, per_channel, channel_axis, n_bits, is_uniform_quantization=True)
-
     if quant_error_method == qc.QuantizationErrorMethod.NOCLIPPING:
-        mm = tensor_min, tensor_max
+        if channel_axis is None and per_channel:
+            total_error_list = []
+            th_list = []
+            for _axis in range(len(tensor_data.shape)):
+                tensor_min = get_tensor_min(tensor_data, per_channel, _axis)
+                tensor_max = get_tensor_max(tensor_data, per_channel, _axis, n_bits, is_uniform_quantization=True)
+                q_tensor_data = uniform_quantize_tensor(tensor_data, tensor_min, tensor_max, n_bits)
+                total_error_list.append(compute_mse(tensor_data, q_tensor_data, norm=True))
+                th_list.append((tensor_min, tensor_max))
+            channel_axis = np.argmin(total_error_list)
+            mm = th_list[channel_axis]
+        else:
+            tensor_min = get_tensor_min(tensor_data, per_channel, channel_axis)
+            tensor_max = get_tensor_max(tensor_data, per_channel, channel_axis, n_bits, is_uniform_quantization=True)
+            mm = tensor_min, tensor_max
     else:
         axis = -1 if per_channel else None
         error_function = get_threshold_selection_tensor_error_function(QuantizationMethod.UNIFORM, quant_error_method,
                                                                        p, axis=axis, norm=False, node=node,
                                                                        hessian_info_service=hessian_info_service,
                                                                        num_hessian_samples=num_hessian_samples)
-        mm = qparams_uniform_selection_tensor_search(error_function,
-                                                     tensor_data,
-                                                     tensor_min,
-                                                     tensor_max,
-                                                     n_bits,
-                                                     per_channel,
-                                                     channel_axis)
+        mm, channel_axis = qparams_uniform_selection_tensor_search(error_function,
+                                                                   tensor_data,
+                                                                   n_bits,
+                                                                   per_channel,
+                                                                   channel_axis)
+    # In case the tensor\axis has a single value, then min==max, so need to adjust either min or max to zero.
+    if not isinstance(mm[0], np.ndarray):
+        if mm[0] > 0:
+            mm = (np.float32(0).astype(mm[0].dtype), mm[1])
+        if mm[1] < 0:
+            mm = (mm[0], np.float32(0).astype(mm[1].dtype))
+    else:
+        adj_min_to_zero = np.logical_and(mm[1] == mm[0], mm[0] > 0)
+        adj_max_to_zero = np.logical_and(mm[1] == mm[0], mm[1] < 0)
+        mm[0][adj_min_to_zero] = 0
+        mm[1][adj_max_to_zero] = 0
     return {RANGE_MIN: mm[0],
-            RANGE_MAX: mm[1]}
+            RANGE_MAX: mm[1]}, channel_axis
 
 
 def uniform_selection_histogram(bins: np.ndarray,

model_compression_toolkit/core/keras/reader/node_builder.py

@@ -158,7 +158,8 @@ def build_node(node: KerasNode,
             if is_const(arg) or (
                     keras_layer.symbol in tf_function_symbols and
                     isinstance(arg, (tuple, list))):
-                weights.update({i: to_numpy(arg, is_single_tensor=True)})
+                if inputs_as_list or i in kwarg2index.values():
+                    weights.update({i: to_numpy(arg, is_single_tensor=True)})
         # remove weights and KerasTensors and weights from op_call_args
         if inputs_as_list:
             op_call_args = tuple(op_call_args[1:])
@@ -169,8 +170,7 @@ def build_node(node: KerasNode,
         # read weights from call kwargs
         weight_keys = []
         for k, v in op_call_kwargs.items():
-            if is_const(v) or (keras_layer.function in [tf.add, tf.multiply, tf.subtract, tf.divide, tf.truediv, tf.pow,
-                                                        tf.matmul] and
+            if is_const(v) or (keras_layer.symbol in tf_function_symbols and
                                isinstance(v, (tuple, list))):
                 if k in kwarg2index:
                     weights.update({kwarg2index[k]: to_numpy(v, is_single_tensor=True)})

model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -39,20 +39,16 @@ from mct_quantizers import PytorchQuantizationWrapper
 def _build_input_tensors_list(node: BaseNode,
                               graph: Graph,
                               inputs: Tuple[Any],
-                              node_to_output_tensors_dict: Dict[BaseNode, List],
-                              is_op_quantize_wrapper: bool) -> List[List]:
+                              node_to_output_tensors_dict: Dict[BaseNode, List]) -> List[List]:
     """
     Given a node, build a list of input tensors the node gets. The list is built based on the
-    node's incoming edges, previous nodes' output tensors and the node's positional weights.
-    Positional weights aren't used if the node's op is PytorchQuantizationWrapper, since it's
-    positional weights are already in the wrapper.
+    node's incoming edges, previous nodes' output tensors.
 
     Args:
         node: Node to build its input tensors list.
         graph: Graph the node is in.
         inputs: list of input tensors to model.
         node_to_output_tensors_dict: A dictionary from a node to its output tensors.
-        is_op_quantize_wrapper: Whether the func_op is a PytorchQuantizationWrapper or not.
 
     Returns:
         A list of the node's input tensors.
@@ -67,35 +63,30 @@ def _build_input_tensors_list(node: BaseNode,
             _input_tensors = node_to_output_tensors_dict[ie.source_node]
             input_tensors.append(_input_tensors)
         input_tensors = [tensor for tensor_list in input_tensors for tensor in tensor_list]  # flat list of lists
-        input_tensors = node.insert_positional_weights_to_input_list(input_tensors)
-        # convert inputs from positional weights (numpy arrays) to tensors. Must handle each element in the
-        # list separately, because in FX the tensors are FX objects and fail to_torch_tensor
-        input_tensors = [to_torch_tensor(t, numpy_type=t.dtype) if isinstance(t, np.ndarray) else t
-                         for t in input_tensors]
     return input_tensors
 
 
 def _merge_inputs(_node: BaseNode, input_tensors: List, op_call_args: List,
-                  is_op_quantize_wrapper: bool) -> List:
+                  tensor_input_indices: List = None) -> List:
     """
-    Merge input tensors list with op_call_args, according to correct order.
+    Merge input tensors list with positional weights and op_call_args, according to correct order.
 
     Args:
         _node: The node the inputs are for.
         input_tensors: activation input tensors to node.
         op_call_args: framework node call args.
-        is_op_quantize_wrapper: Whether the func_op is a PytorchQuantizationWrapper or not.
+
     Returns:
         Combined list of input_tensors and op_call_args.
     """
     if isinstance(_node, FunctionalNode) and _node.tensor_input_indices:
         _input_list = op_call_args.copy()
-        if is_op_quantize_wrapper:
-            _input_list = input_tensors + _input_list
-        else:
-            assert len(_node.tensor_input_indices) == len(input_tensors), 'Mismatch between input tensors and indices'
-            for i, t in zip(_node.tensor_input_indices, input_tensors):
-                _input_list.insert(i, t)
+        if tensor_input_indices is None:
+            tensor_input_indices = _node.tensor_input_indices
+        assert len(tensor_input_indices) == len(input_tensors), \
+            f'Mismatch between input tensors ({len(tensor_input_indices)}) and indices {len(input_tensors)}'
+        for i, t in zip(tensor_input_indices, input_tensors):
+            _input_list.insert(i, t)
     else:
         _input_list = input_tensors + op_call_args
 
@@ -126,10 +117,22 @@ def _run_operation(n: BaseNode,
     op_call_args = n.op_call_args if isinstance(n, FunctionalNode) else []
     functional_kwargs = n.op_call_kwargs if isinstance(n, FunctionalNode) else {}
 
+    if not (isinstance(n, FunctionalNode) and isinstance(op_func, PytorchQuantizationWrapper)):
+        # Insert positional weights only when not a quantized functional node, because quantized functional nodes
+        # insert the quantized weights in the wrapper.
+        input_tensors = n.insert_positional_weights_to_input_list(input_tensors)
+        # convert inputs from positional weights (numpy arrays) to tensors. Must handle each element in the
+        # list separately, because in FX the tensors are FX objects and fail to_torch_tensor
+        input_tensors = [to_torch_tensor(t, numpy_type=t.dtype) if isinstance(t, np.ndarray) else t
+                         for t in input_tensors]
+        _tensor_input_indices = None
+    else:
+        _tensor_input_indices = [i for i in n.tensor_input_indices if i not in n.weights]
+
     if isinstance(n, FunctionalNode) and n.inputs_as_list:
         out_tensors_of_n_float = op_func(input_tensors, *op_call_args, **functional_kwargs)
     else:
-        merged_inputs = _merge_inputs(n, input_tensors, op_call_args, isinstance(op_func, PytorchQuantizationWrapper))
+        merged_inputs = _merge_inputs(n, input_tensors, op_call_args, tensor_input_indices=_tensor_input_indices)
         out_tensors_of_n_float = op_func(*merged_inputs, **functional_kwargs)
 
     # Add a fake quant node if the node has an activation threshold.
@@ -295,8 +298,7 @@ class PytorchModel(torch.nn.Module):
             input_tensors = _build_input_tensors_list(node,
                                                       self.graph,
                                                       args,
-                                                      node_to_output_tensors_dict,
-                                                      isinstance(op_func, PytorchQuantizationWrapper))
+                                                      node_to_output_tensors_dict)
             use_activation_quantization, activation_quantization_fn = self._get_activation_quantization_fn(node)
 
             # Run node operation and fetch outputs