mct-nightly 2.1.0.20240811.503__py3-none-any.whl → 2.1.0.20240813.442__py3-none-any.whl

model_compression_toolkit/core/common/quantization/set_node_quantization_config.py

@@ -15,9 +15,11 @@
 
 
 import copy
-from typing import List, Tuple
+from typing import List, Tuple,  Optional
 
+from mct_quantizers.common.constants import ACTIVATION_N_BITS
 from model_compression_toolkit.core.common import BaseNode
+from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.graph.base_graph import Graph
@@ -37,19 +39,21 @@ from model_compression_toolkit.target_platform_capabilities.target_platform.op_q
 
 def set_quantization_configuration_to_graph(graph: Graph,
                                             quant_config: QuantizationConfig,
+                                            bit_width_config: BitWidthConfig = None,
                                             mixed_precision_enable: bool = False,
                                             running_gptq: bool = False) -> Graph:
     """
     Add quantization configuration for each graph node.
 
     Args:
-        graph: Graph for which to add quantization info to each node.
-        quant_config: Quantization configuration containing parameters for how the graph should be quantized.
-        mixed_precision_enable: is mixed precision enabled.
-        running_gptq: Whether or not a GPTQ optimization is planned to run after the PTQ process.
+        graph (Graph): Graph for which to add quantization info to each node.
+        quant_config (QuantizationConfig): Quantization configuration containing parameters for how the graph should be quantized.
+        bit_width_config (BitWidthConfig): Configuration for manual bit width selection. Defaults to None.
+        mixed_precision_enable (bool): Whether mixed precision is enabled. Defaults to False.
+        running_gptq (bool): Whether or not a GPTQ optimization is planned to run after the PTQ process. Defaults to False.
 
     Returns:
-        The graph with quantization configurations attached to each node in it.
+        Graph: The graph with quantization configurations attached to each node in it.
     """
 
     if quant_config.weights_error_method == QuantizationErrorMethod.HMSE:
@@ -62,13 +66,16 @@ def set_quantization_configuration_to_graph(graph: Graph,
             Logger.warning("Using the HMSE error method for weights quantization parameters search. "
                            "Note: This method may significantly increase runtime during the parameter search process.")
 
+    nodes_to_manipulate_bit_widths = {} if bit_width_config is None else bit_width_config.get_nodes_to_manipulate_bit_widths(graph)
+
     for n in graph.nodes:
         set_quantization_configs_to_node(node=n,
                                          graph=graph,
                                          quant_config=quant_config,
                                          fw_info=graph.fw_info,
                                          tpc=graph.tpc,
-                                         mixed_precision_enable=mixed_precision_enable)
+                                         mixed_precision_enable=mixed_precision_enable,
+                                         manual_bit_width_override=nodes_to_manipulate_bit_widths.get(n))
     return graph
 
 
@@ -77,21 +84,32 @@ def set_quantization_configs_to_node(node: BaseNode,
                                      quant_config: QuantizationConfig,
                                      fw_info: FrameworkInfo,
                                      tpc: TargetPlatformCapabilities,
-                                     mixed_precision_enable: bool = False):
+                                     mixed_precision_enable: bool = False,
+                                     manual_bit_width_override: Optional[int] = None):
     """
     Create and set quantization configurations to a node (for both weights and activation).
 
     Args:
-        node: Node to set its quantization configurations.
-        graph: Model's internal representation graph.
-        quant_config: Quantization configuration to generate the node's configurations from.
-        fw_info: Information needed for quantization about the specific framework.
-        tpc: TargetPlatformCapabilities to get default OpQuantizationConfig.
-        mixed_precision_enable: is mixed precision enabled.
+        node (BaseNode): Node to set its quantization configurations.
+        graph (Graph): Model's internal representation graph.
+        quant_config (QuantizationConfig): Quantization configuration to generate the node's configurations from.
+        fw_info (FrameworkInfo): Information needed for quantization about the specific framework.
+        tpc (TargetPlatformCapabilities): TargetPlatformCapabilities to get default OpQuantizationConfig.
+        mixed_precision_enable (bool): Whether mixed precision is enabled. Defaults to False.
+        manual_bit_width_override (Optional[int]): Specifies a custom bit-width to override the node's activation bit-width. Defaults to None.
     """
     node_qc_options = node.get_qco(tpc)
     base_config, node_qc_options_list = node.filter_node_qco_by_graph(tpc, graph.get_next_nodes(node), node_qc_options)
 
+    # If a manual_bit_width_override is given, filter node_qc_options_list to retain only the options with activation bits equal to manual_bit_width_override,
+    # and update base_config accordingly.
+    base_config, node_qc_options_list = filter_qc_options_with_manual_bit_width(
+        node=node,
+        node_qc_options_list=node_qc_options_list,
+        base_config=base_config,
+        manual_bit_width_override=manual_bit_width_override,
+        mixed_precision_enable=mixed_precision_enable)
+
     # Create QC candidates for weights and activation combined
     weight_channel_axis = fw_info.kernel_channels_mapping.get(node.type)
     node.candidates_quantization_cfg = _create_node_candidates_qc(quant_config,
@@ -199,16 +217,16 @@ def _create_node_candidates_qc(qc: QuantizationConfig,
     Create a list of candidates of weights and activation quantization configurations for a node.
 
     Args:
-        qc: Quantization configuration the quantization process should follow.
-        fw_info: Framework information (e.g., which layers should have their kernels' quantized).
-        weight_channel_axis: (Output, Input) channel index of the node's kernel.
-        node_qc_options_list: List of quantization configs of node.
-        base_config: Base quantization config for node.
-        node: A node to set quantization configuration candidates to.
-        mixed_precision_enable: is mixed precision enabled
+        qc (QuantizationConfig): Quantization configuration the quantization process should follow.
+        fw_info (FrameworkInfo): Framework information (e.g., which layers should have their kernels quantized).
+        weight_channel_axis (Tuple[int, int]): (Output, Input) channel index of the node's kernel.
+        node_qc_options_list (List[OpQuantizationConfig]): List of quantization configs of node.
+        base_config (OpQuantizationConfig): Base quantization config for node.
+        node (BaseNode): A node to set quantization configuration candidates to.
+        mixed_precision_enable (bool): Whether mixed precision is enabled. Defaults to False.
 
     Returns:
-        List of candidates of weights quantization configurations to set for a node.
+        List[CandidateNodeQuantizationConfig]: List of candidates of weights quantization configurations to set for a node.
     """
 
     candidates = []
@@ -231,3 +249,51 @@ def _create_node_candidates_qc(qc: QuantizationConfig,
                                                            node_attrs_list))
 
     return candidates
+
+
+def filter_qc_options_with_manual_bit_width(
+        node: BaseNode,
+        node_qc_options_list: List[OpQuantizationConfig],
+        base_config: OpQuantizationConfig,
+        manual_bit_width_override: Optional[int],
+        mixed_precision_enable: bool) -> Tuple[OpQuantizationConfig, List[OpQuantizationConfig]]:
+    """
+    Update the quantization configurations for a node, allowing manual bit-width overrides if specified.
+
+    Args:
+        node (BaseNode): A node to set quantization configuration candidates to.
+        node_qc_options_list (List[OpQuantizationConfig]): List of quantization configs for the node.
+        base_config (OpQuantizationConfig): Base quantization config for the node.
+        manual_bit_width_override (Optional[int]): Specifies a custom bit-width to override the node's activation bit-width.
+        mixed_precision_enable (bool): Whether mixed precision is enabled.
+
+    Returns:
+        Tuple[OpQuantizationConfig, List[OpQuantizationConfig]]: The updated base configuration and the filtered list of quantization configs.
+    """
+    if manual_bit_width_override is None:
+        return base_config, node_qc_options_list
+
+    # Filter node_qc_options_list to retain only the options with activation bits equal to manual_bit_width_override.
+    node_qc_options_list = [op_cfg for op_cfg in node_qc_options_list if
+                                manual_bit_width_override == op_cfg.activation_n_bits]
+
+    if len(node_qc_options_list) == 0:
+        Logger.critical(f"Manually selected activation bit-width {manual_bit_width_override} is invalid for node {node}.")
+    else:
+        # Update the base_config to one of the values from the filtered node_qc_options_list.
+        # First, check if a configuration similar to the original base_config but with activation bits equal to manual_bit_width_override exists.
+        # If it does, use it as the base_config. If not, choose a different configuration from node_qc_options_list.
+        Logger.info(f"Setting node {node} bit-width to manually selected bit-width: {manual_bit_width_override} bits.")
+        updated_base_config = base_config.clone_and_edit({ACTIVATION_N_BITS, manual_bit_width_override})
+        if updated_base_config in node_qc_options_list:
+            # If a base_config with the specified manual_bit_width_override exists in the node_qc_options_list,
+            # point the base_config to this option.
+            base_config = node_qc_options_list[node_qc_options_list.index(updated_base_config)]
+        else:
+            # Choose a different configuration from node_qc_options_list. If multiple options exist, issue a warning.
+            base_config = node_qc_options_list[0]
+            if len(node_qc_options_list) > 0 and not mixed_precision_enable:
+                Logger.info(
+                    f"Request received to select {manual_bit_width_override} activation bits. However, the base configuration for layer type {node.type} is missing in the node_qc_options_list."
+                    f" Overriding base_config with an option that uses {manual_bit_width_override} bit activations.")  # pragma: no cover
+    return base_config, node_qc_options_list

model_compression_toolkit/core/graph_prep_runner.py

@@ -20,6 +20,7 @@ from model_compression_toolkit.core.common import FrameworkInfo
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.fusion.layer_fusing import fusion
 from model_compression_toolkit.core.common.graph.base_graph import Graph
+from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.filter_nodes_candidates import filter_nodes_candidates
 from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
@@ -38,6 +39,7 @@ def graph_preparation_runner(in_model: Any,
                              fw_info: FrameworkInfo,
                              fw_impl: FrameworkImplementation,
                              tpc: TargetPlatformCapabilities,
+                             bit_width_config: BitWidthConfig = None,
                              tb_w: TensorboardWriter = None,
                              mixed_precision_enable: bool = False,
                              running_gptq: bool = False) -> Graph:
@@ -50,17 +52,18 @@ def graph_preparation_runner(in_model: Any,
         - Apply all necessary substitutions to finalize the graph for quantization.
 
     Args:
-        in_model: Model to quantize.
-        representative_data_gen: Dataset used for calibration.
-        quantization_config: QuantizationConfig containing parameters of how the model should be quantized.
-        fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
+        in_model (Any): Model to quantize.
+        representative_data_gen (Callable): Dataset used for calibration.
+        quantization_config (QuantizationConfig): QuantizationConfig containing parameters of how the model should be quantized.
+        fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.).
-        fw_impl: FrameworkImplementation object with a specific framework methods implementation.
-        tpc: TargetPlatformCapabilities object that models the inference target platform and
+        fw_impl (FrameworkImplementation): FrameworkImplementation object with a specific framework methods implementation.
+        tpc (TargetPlatformCapabilities): TargetPlatformCapabilities object that models the inference target platform and
             the attached framework operator's information.
-        tb_w: TensorboardWriter object for logging.
-        mixed_precision_enable: is mixed precision enabled.
-        running_gptq: Whether or not a GPTQ optimization is planned to run after the PTQ process.
+        bit_width_config (BitWidthConfig): Config for bit-width selection. Defaults to None.
+        tb_w (TensorboardWriter): TensorboardWriter object for logging.
+        mixed_precision_enable (bool): is mixed precision enabled.
+        running_gptq (bool): Whether or not a GPTQ optimization is planned to run after the PTQ process.
 
     Returns:
         An internal graph representation of the input model.
@@ -78,6 +81,7 @@ def graph_preparation_runner(in_model: Any,
     transformed_graph = get_finalized_graph(graph,
                                             tpc,
                                             quantization_config,
+                                            bit_width_config,
                                             fw_info,
                                             tb_w,
                                             fw_impl,
@@ -90,6 +94,7 @@ def graph_preparation_runner(in_model: Any,
 def get_finalized_graph(initial_graph: Graph,
                         tpc: TargetPlatformCapabilities,
                         quant_config: QuantizationConfig = DEFAULTCONFIG,
+                        bit_width_config: BitWidthConfig = None,
                         fw_info: FrameworkInfo = None,
                         tb_w: TensorboardWriter = None,
                         fw_impl: FrameworkImplementation = None,
@@ -104,6 +109,7 @@ def get_finalized_graph(initial_graph: Graph,
         tpc (TargetPlatformCapabilities): TargetPlatformCapabilities object that describes the desired inference target platform (includes fusing patterns MCT should handle).
         quant_config (QuantizationConfig): QuantizationConfig containing parameters of how the model should be
             quantized.
+        bit_width_config (BitWidthConfig): Config for bit-width selection. Defaults to None.
         fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g.,
             kernel channels indices, groups of layers by how they should be quantized, etc.)
         tb_w (TensorboardWriter): TensorboardWriter object to use for logging events such as graphs, histograms, etc.
@@ -147,6 +153,7 @@ def get_finalized_graph(initial_graph: Graph,
     ######################################
     transformed_graph = set_quantization_configuration_to_graph(graph=transformed_graph,
                                                                 quant_config=quant_config,
+                                                                bit_width_config=bit_width_config,
                                                                 mixed_precision_enable=mixed_precision_enable,
                                                                 running_gptq=running_gptq)
 

model_compression_toolkit/core/keras/resource_utilization_data_facade.py

@@ -20,7 +20,7 @@ from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW
 from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization_data import compute_resource_utilization_data
-from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.verify_packages import FOUND_TF
 
 if FOUND_TF:
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
@@ -89,5 +89,6 @@ else:
     # If tensorflow is not installed,
     # we raise an exception when trying to use this function.
     def keras_resource_utilization_data(*args, **kwargs):
-        Logger.critical("Tensorflow must be installed to use keras_resource_utilization_data. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "keras_resource_utilization_data. The 'tensorflow' package is either not installed or is "
+                        "installed with a version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/core/pytorch/resource_utilization_data_facade.py

@@ -23,7 +23,7 @@ from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization_data import compute_resource_utilization_data
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
-from model_compression_toolkit.constants import FOUND_TORCH
+from model_compression_toolkit.verify_packages import FOUND_TORCH
 
 if FOUND_TORCH:
     from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO

model_compression_toolkit/core/runner.py

@@ -115,6 +115,7 @@ def core_runner(in_model: Any,
                                      fw_info,
                                      fw_impl,
                                      tpc,
+                                     core_config.bit_width_config,
                                      tb_w,
                                      mixed_precision_enable=core_config.mixed_precision_enable,
                                      running_gptq=running_gptq)

model_compression_toolkit/data_generation/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.constants import FOUND_TORCH, FOUND_TF, FOUND_TORCHVISION
+from model_compression_toolkit.verify_packages import FOUND_TORCHVISION, FOUND_TORCH, FOUND_TF
 from model_compression_toolkit.data_generation.common.data_generation_config import DataGenerationConfig
 from model_compression_toolkit.data_generation.common.enums import ImageGranularity, DataInitType, SchedulerType, BNLayerWeightingType, OutputLossType, BatchNormAlignemntLossType, ImagePipelineType, ImageNormalizationType
 

model_compression_toolkit/data_generation/keras/keras_data_generation.py

@@ -16,7 +16,7 @@ import time
 from typing import Callable, Tuple, List, Dict, Union
 from tqdm import tqdm
 
-from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.data_generation.common.constants import DEFAULT_N_ITER, DEFAULT_DATA_GEN_BS
 from model_compression_toolkit.data_generation.common.data_generation import get_data_generation_classes
 from model_compression_toolkit.data_generation.common.image_pipeline import image_normalization_dict
@@ -349,8 +349,12 @@ if FOUND_TF:
 else:
     def get_keras_data_generation_config(*args, **kwargs):
         Logger.critical(
-            "Tensorflow must be installed to use get_tensorflow_data_generation_config. The 'tensorflow' package is missing.")  # pragma: no cover
+            "Tensorflow must be installed with a version of 2.15 or lower to use "
+            "get_tensorflow_data_generation_config. The 'tensorflow' package is missing or is installed with a "
+            "version higher than 2.15.")  # pragma: no cover
 
 
     def keras_data_generation_experimental(*args, **kwargs):
-        Logger.critical("Tensorflow must be installed to use tensorflow_data_generation_experimental. The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "tensorflow_data_generation_experimental. The 'tensorflow' package is missing or is installed "
+                        "with a version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py

@@ -17,7 +17,7 @@ from typing import Callable, Any, Tuple, List, Union
 
 from tqdm import tqdm
 
-from model_compression_toolkit.constants import FOUND_TORCH, FOUND_TORCHVISION
+from model_compression_toolkit.verify_packages import FOUND_TORCHVISION, FOUND_TORCH
 from model_compression_toolkit.core.pytorch.utils import set_model
 from model_compression_toolkit.data_generation.common.constants import DEFAULT_N_ITER, DEFAULT_DATA_GEN_BS
 from model_compression_toolkit.data_generation.common.data_generation import get_data_generation_classes

model_compression_toolkit/exporter/model_exporter/keras/keras_export_facade.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from typing import Callable, Dict
 
-from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.exporter.model_exporter.fw_agonstic.quantization_format import QuantizationFormat
 from model_compression_toolkit.logger import Logger
 
@@ -101,5 +101,6 @@ if FOUND_TF:
         return exporter.get_custom_objects()
 else:
     def keras_export_model(*args, **kwargs):
-        Logger.critical("Tensorflow must be installed to use keras_export_model. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use keras_export_model."
+                        "The 'tensorflow' package is missing or is installed "
+                        "with a version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -18,7 +18,7 @@ from io import BytesIO
 import torch.nn
 
 from mct_quantizers import PytorchActivationQuantizationHolder, PytorchQuantizationWrapper
-from model_compression_toolkit.constants import FOUND_ONNX
+from model_compression_toolkit.verify_packages import FOUND_ONNX
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from model_compression_toolkit.exporter.model_exporter.pytorch.base_pytorch_exporter import BasePyTorchExporter

model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from typing import Callable
 
-from model_compression_toolkit.constants import FOUND_TORCH
+from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.exporter.model_exporter.fw_agonstic.quantization_format import QuantizationFormat
 from model_compression_toolkit.exporter.model_exporter.pytorch.export_serialization_format import \
     PytorchExportSerializationFormat

model_compression_toolkit/exporter/model_wrapper/keras/builder/fully_quantized_model_builder.py

@@ -16,7 +16,7 @@
 from typing import Tuple, Callable
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.logger import Logger
 import model_compression_toolkit.core as C
@@ -101,5 +101,6 @@ if FOUND_TF:
         return exportable_model, user_info
 else:
     def get_exportable_keras_model(*args, **kwargs):  # pragma: no cover
-        Logger.critical("Tensorflow must be installed to use get_exportable_keras_model. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "get_exportable_keras_model. The 'tensorflow' package is missing or is installed with a "
+                        "version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/exporter/model_wrapper/keras/validate_layer.py

@@ -15,7 +15,7 @@
 from typing import Any
 
 from mct_quantizers import BaseInferableQuantizer, KerasActivationQuantizationHolder
-from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.logger import Logger
 
 if FOUND_TF:
@@ -76,5 +76,6 @@ if FOUND_TF:
         return True
 else:
     def is_keras_layer_exportable(*args, **kwargs):  # pragma: no cover
-        Logger.critical("Tensorflow must be installed to use is_keras_layer_exportable. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "is_keras_layer_exportable. The 'tensorflow' package is missing or is installed with a "
+                        "version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py

@@ -16,7 +16,7 @@
 from typing import Union, Callable
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph
-from model_compression_toolkit.constants import FOUND_TORCH
+from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common import BaseNode
 import model_compression_toolkit.core as C

model_compression_toolkit/exporter/model_wrapper/pytorch/validate_layer.py

@@ -15,8 +15,7 @@
 from typing import Any
 
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.constants import FOUND_TORCH
-
+from model_compression_toolkit.verify_packages import FOUND_TORCH
 
 if FOUND_TORCH:
     import torch.nn as nn

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -21,7 +21,8 @@ from model_compression_toolkit.core.common.quantization.quantize_graph_weights i
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.constants import TENSORFLOW, FOUND_TF, ACT_HESSIAN_DEFAULT_BATCH_SIZE
+from model_compression_toolkit.constants import TENSORFLOW, ACT_HESSIAN_DEFAULT_BATCH_SIZE
+from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, GPTQHessianScoresConfig
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
@@ -251,10 +252,12 @@ else:
     # If tensorflow is not installed,
     # we raise an exception when trying to use these functions.
     def get_keras_gptq_config(*args, **kwargs):
-        Logger.critical("Tensorflow must be installed to use get_keras_gptq_config. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "get_keras_gptq_config. The 'tensorflow' package is missing or is "
+                        "installed with a version higher than 2.15.")  # pragma: no cover
 
 
     def keras_gradient_post_training_quantization(*args, **kwargs):
-        Logger.critical("Tensorflow must be installed to use keras_gradient_post_training_quantization. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "keras_gradient_post_training_quantization. The 'tensorflow' package is missing or is "
+                        "installed with a version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py

@@ -16,7 +16,7 @@ from abc import abstractmethod
 from typing import Union, Dict, List
 
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.gptq.common.gptq_constants import WEIGHTS_QUANTIZATION_PARAMS
 
 from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
@@ -105,5 +105,6 @@ if FOUND_TF:
 else:
     class BaseKerasGPTQTrainableQuantizer:  # pragma: no cover
         def __init__(self, *args, **kwargs):
-            Logger.critical("Tensorflow must be installed to use BaseKerasGPTQTrainableQuantizer. "
-                            "The 'tensorflow' package is missing.")  # pragma: no cover
+            Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                            "BaseKerasGPTQTrainableQuantizer. The 'tensorflow' package is missing or is "
+                            "installed with a version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -16,7 +16,8 @@ import copy
 
 from typing import Callable
 from model_compression_toolkit.core import common
-from model_compression_toolkit.constants import FOUND_TORCH, ACT_HESSIAN_DEFAULT_BATCH_SIZE
+from model_compression_toolkit.constants import ACT_HESSIAN_DEFAULT_BATCH_SIZE
+from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
 from model_compression_toolkit.logger import Logger

model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py

@@ -16,7 +16,7 @@ from abc import abstractmethod
 from typing import Union, Dict
 
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.constants import FOUND_TORCH
+from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.gptq.common.gptq_constants import WEIGHTS_QUANTIZATION_PARAMS
 
 from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, \

model_compression_toolkit/pruning/keras/pruning_facade.py

@@ -16,11 +16,13 @@
 from typing import Callable, Tuple
 
 from model_compression_toolkit import get_target_platform_capabilities
-from model_compression_toolkit.constants import TENSORFLOW, FOUND_TF
+from model_compression_toolkit.constants import TENSORFLOW
+from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.pruning.pruner import Pruner
 from model_compression_toolkit.core.common.pruning.pruning_config import PruningConfig
 from model_compression_toolkit.core.common.pruning.pruning_info import PruningInfo
+from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.set_node_quantization_config import set_quantization_configuration_to_graph
 from model_compression_toolkit.core.graph_prep_runner import read_model_to_graph
 from model_compression_toolkit.logger import Logger
@@ -148,5 +150,6 @@ else:
     # If tensorflow is not installed,
     # we raise an exception when trying to use these functions.
     def keras_pruning_experimental(*args, **kwargs):
-        Logger.critical("Tensorflow must be installed to use keras_pruning_experimental. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "keras_pruning_experimental. The 'tensorflow' package is missing or is "
+                        "installed with a version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/pruning/pytorch/pruning_facade.py

@@ -15,11 +15,13 @@
 
 from typing import Callable, Tuple
 from model_compression_toolkit import get_target_platform_capabilities
-from model_compression_toolkit.constants import FOUND_TORCH, PYTORCH
+from model_compression_toolkit.constants import PYTORCH
+from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.pruning.pruner import Pruner
 from model_compression_toolkit.core.common.pruning.pruning_config import PruningConfig
 from model_compression_toolkit.core.common.pruning.pruning_info import PruningInfo
+from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.set_node_quantization_config import set_quantization_configuration_to_graph
 from model_compression_toolkit.core.graph_prep_runner import read_model_to_graph
 from model_compression_toolkit.logger import Logger

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -21,7 +21,8 @@ from model_compression_toolkit.core.analyzer import analyzer_model_quantization
 from model_compression_toolkit.core.common.quantization.quantize_graph_weights import quantize_graph_weights
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.constants import TENSORFLOW, FOUND_TF
+from model_compression_toolkit.constants import TENSORFLOW
+from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfig
@@ -178,5 +179,6 @@ else:
     # If tensorflow is not installed,
     # we raise an exception when trying to use these functions.
     def keras_post_training_quantization(*args, **kwargs):
-        Logger.critical("Tensorflow must be installed to use keras_post_training_quantization. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "keras_post_training_quantization. The 'tensorflow' package is missing or is "
+                        "installed with a version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -18,7 +18,8 @@ from typing import Callable
 
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.constants import PYTORCH, FOUND_TORCH
+from model_compression_toolkit.constants import PYTORCH
+from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core import CoreConfig

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -19,7 +19,7 @@ from functools import partial
 from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfig
@@ -291,10 +291,12 @@ else:
     # If tensorflow is not installed,
     # we raise an exception when trying to use these functions.
     def keras_quantization_aware_training_init_experimental(*args, **kwargs):
-        Logger.critical("Tensorflow must be installed to use keras_quantization_aware_training_init_experimental. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "keras_quantization_aware_training_init_experimental. The 'tensorflow' package is missing "
+                        "or is installed with a version higher than 2.15.")  # pragma: no cover
 
 
     def keras_quantization_aware_training_finalize_experimental(*args, **kwargs):
-        Logger.critical("Tensorflow must be installed to use keras_quantization_aware_training_finalize_experimental. "
-                        "The 'tensorflow' package is missing.")  # pragma: no cover
+        Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                        "keras_quantization_aware_training_finalize_experimental. The 'tensorflow' package is missing "
+                        "or is installed with a version higher than 2.15.")  # pragma: no cover

model_compression_toolkit/qat/keras/quantizer/base_keras_qat_quantizer.py

@@ -15,7 +15,7 @@
 from typing import Union
 
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.verify_packages import FOUND_TF
 
 from model_compression_toolkit.trainable_infrastructure import TrainableQuantizerWeightsConfig, \
     TrainableQuantizerActivationConfig, BaseKerasTrainableQuantizer
@@ -44,5 +44,6 @@ else:  # pragma: no cover
                      quantization_config: Union[TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig]):
 
             super().__init__(quantization_config)
-            Logger.critical("Tensorflow must be installed to use BaseKerasQATTrainableQuantizer. "
-                            "The 'tensorflow' package is missing.")  # pragma: no cover
+            Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
+                            "BaseKerasQATTrainableQuantizer. The 'tensorflow' package is missing "
+                            "or is installed with a version higher than 2.15.")  # pragma: no cover