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

{mct_nightly-2.1.0.20240811.503.dist-info → mct_nightly-2.1.0.20240812.432.dist-info}/METADATA

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

{mct_nightly-2.1.0.20240811.503.dist-info → mct_nightly-2.1.0.20240812.432.dist-info}/RECORD

@@ -1,13 +1,13 @@
-model_compression_toolkit/__init__.py,sha256=28oZ77PKcaaEtbwP3jR4kM5WSpNYQXslUm933JqpuUw,1573
+model_compression_toolkit/__init__.py,sha256=dFfNYHNevEMx3n6CPcaXHcIQxE4Nlhrsckn2CtIDLiY,1573
 model_compression_toolkit/constants.py,sha256=0qrEGjX36Oo7Lt8mR0LD2aSe2xA7gKrhkzBGp7g5eiA,4345
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
 model_compression_toolkit/metadata.py,sha256=UtXS5ClK-qPoxGRuytlDGZSzgLo911dMni2EFRcg6io,3623
-model_compression_toolkit/core/__init__.py,sha256=TrRgkWpT1AN2Faw1M_1HXyJkJnbxfn9p-RigDZl7pg0,1982
+model_compression_toolkit/core/__init__.py,sha256=tnDtL9KmT0vsOU27SsJ19TKDEbIH-tXYeGxTo5YnNUM,2077
 model_compression_toolkit/core/analyzer.py,sha256=X-2ZpkH1xdXnISnw1yJvXnvV-ssoUh-9LkLISSWNqiY,3691
-model_compression_toolkit/core/graph_prep_runner.py,sha256=kM70wmNG3yMFiGQc0uO0wn9j4ZbSWxUEykpxDK55doc,10567
+model_compression_toolkit/core/graph_prep_runner.py,sha256=7-b7Jd5jBVaXOWg5nSqbEyzBtdaGDbCxs8aqMV6GZ6I,11287
 model_compression_toolkit/core/quantization_prep_runner.py,sha256=K9eJ7VbB_rpeyxX4yEnorOmSxFW3DkvofzxS6QI8Hp8,6454
-model_compression_toolkit/core/runner.py,sha256=uXpyYaX1uFNhKituGmSfKb3ZkguXG2V_Cg6XCnprplg,13569
+model_compression_toolkit/core/runner.py,sha256=XQDNJirZkVJ_FXP72d7tbVc_Tr3Jw0Eqm_kxNHW8kPs,13636
 model_compression_toolkit/core/common/__init__.py,sha256=Wh127PbXcETZX_d1PQqZ71ETK3J9XO5A-HpadGUbj6o,1447
 model_compression_toolkit/core/common/base_substitutions.py,sha256=xDFSmVVs_iFSZfajytI0cuQaNRNcwHX3uqOoHgVUvxQ,1666
 model_compression_toolkit/core/common/framework_implementation.py,sha256=kSg2f7wS7e2EyvX6y0eKfNTTFvVFVrB8lvldJvcPvN8,20724
@@ -70,7 +70,7 @@ model_compression_toolkit/core/common/mixed_precision/set_layer_to_bitwidth.py,s
 model_compression_toolkit/core/common/mixed_precision/solution_refinement_procedure.py,sha256=KifDMbm7qkSfvSl6pcZzQ82naIXzeKL6aT-VsvWZYyc,7901
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization.py,sha256=HILF7CIn-GYPvPmTFyvjWLhuLDwSGwdBcAaKFgVYrwk,4745
-model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_data.py,sha256=az0XfBPVm1kAfxNCPb0Z-Q05-F-vqnmyRpKm6SBLa6c,13826
+model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_data.py,sha256=3ZOI-RNp5faT-U2Og7rLW9EKwBB6ooa7-RwSsWJmquo,14022
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/ru_aggregation_methods.py,sha256=ttc8wPa_9LZansutQ2f1ss-RTzgTv739wy3qsdLzyyk,4217
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/ru_functions_mapping.py,sha256=QhuqaECEGLnYC08iD6-2XXcU7NXbPzYf1sQcjYlGak8,1682
 model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/ru_methods.py,sha256=WC1EHoNuo_lrzy4NRhGJ1cgmJ2IsFsbmP86mrVO3AVA,21506
@@ -98,8 +98,9 @@ model_compression_toolkit/core/common/pruning/mask/__init__.py,sha256=huHoBUcKNB
 model_compression_toolkit/core/common/pruning/mask/per_channel_mask.py,sha256=APY8BsM9B7ZxVCH6n1xs9fSCTB_A9ou9gHrCQl1DOdI,5131
 model_compression_toolkit/core/common/pruning/mask/per_simd_group_mask.py,sha256=4ohJrJHNzZk5uMnZEYkwLx2TDGzkh5kRhLGNVYNC6dc,5978
 model_compression_toolkit/core/common/quantization/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
+model_compression_toolkit/core/common/quantization/bit_width_config.py,sha256=r6VQXgyJxX_AM1JTzv-sTcrvCTnktBfOkVP20RllNmk,4586
 model_compression_toolkit/core/common/quantization/candidate_node_quantization_config.py,sha256=yU-Cr6S4wOSkDk57iH2NVe-WII0whOhLryejkomCOt4,4940
-model_compression_toolkit/core/common/quantization/core_config.py,sha256=KYdyfSmjSL4ye24nKlC_c4_AxYb14qoqaeMnZj4-8kE,2257
+model_compression_toolkit/core/common/quantization/core_config.py,sha256=f0uSuY9mX-vLX_1s2DemPARQlAXmLPKJKPtCArz3pZI,2670
 model_compression_toolkit/core/common/quantization/debug_config.py,sha256=8G8SpE_4rb8xBp8d6mMq8R_OnXJ_1oxB2g-Lxk9EJCM,1691
 model_compression_toolkit/core/common/quantization/filter_nodes_candidates.py,sha256=fwF4VILaX-u3ZaFd81xjbJuhg8Ef-JX_KfMXW0TPV-I,7136
 model_compression_toolkit/core/common/quantization/node_quantization_config.py,sha256=YycYN8_JMzvSR3pTVm5dT5x4zP3yBHn0Z9agnwrvOKI,26395
@@ -108,7 +109,7 @@ model_compression_toolkit/core/common/quantization/quantization_fn_selection.py,
 model_compression_toolkit/core/common/quantization/quantization_params_fn_selection.py,sha256=MwIOBZ4BlZSTIOG75PDvlI3JmZ6t8YjPc1VP9Adei60,3847
 model_compression_toolkit/core/common/quantization/quantize_graph_weights.py,sha256=N005MSvx8UypVpa7XrxNrB2G732n2wHj3RmLyjTgd3I,2728
 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=9Y4eVDWCXFvCaXy2gbb-1880sp52M8wqH0M3KgAw8rM,12834
+model_compression_toolkit/core/common/quantization/set_node_quantization_config.py,sha256=0pZVO4wsNP815R9ZOd5ojC_OdNEeKkxYKdjggsqsZKg,17750
 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=Fd_gxr5js-mqEwucaRR1CQAZ1W_wna19L1gAPeOzxRQ,23610
 model_compression_toolkit/core/common/quantization/quantization_params_generation/lut_kmeans_params.py,sha256=RL-PklAjGyC-26anSt8fU07a6pB_LBQFQy9o4e9giN0,8739
@@ -380,9 +381,9 @@ model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/__init__.py,sha256
 model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/symmetric_ste.py,sha256=6uxq_w62jn8DDOt9T7VtA6jZ8jTAPcbTufKFOYpVUm4,8768
 model_compression_toolkit/pruning/__init__.py,sha256=lQMZS8G0pvR1LVi53nnJHNXgLNTan_MWMdwsVxhjrow,1106
 model_compression_toolkit/pruning/keras/__init__.py,sha256=3Lkr37Exk9u8811hw8hVqkGcbTQGcLjd3LLuLC3fa_E,698
-model_compression_toolkit/pruning/keras/pruning_facade.py,sha256=Vt9ipysniwQw4erWhwMO4oMCpIFUMKIGq67ugieMZd8,8612
+model_compression_toolkit/pruning/keras/pruning_facade.py,sha256=vDpY97xirGF-o5XB6HvG_y2bL4LzfiTW3cPURTvaeKI,8707
 model_compression_toolkit/pruning/pytorch/__init__.py,sha256=pKAdbTCFM_2BrZXUtTIw0ouKotrWwUDF_hP3rPwCM2k,696
-model_compression_toolkit/pruning/pytorch/pruning_facade.py,sha256=cSuvHHCqgr7k9FdYOxFqe2njLcJ7IkzCrWSb26S0TK8,9398
+model_compression_toolkit/pruning/pytorch/pruning_facade.py,sha256=1uo5jWgbFNNhRbfb8da5REymMUdLJ3JidR8aAMXCBoE,9493
 model_compression_toolkit/ptq/__init__.py,sha256=Z_hkmTh7aLFei1DJKV0oNVUbrv_Q_0CTw-qD85Xf8UM,904
 model_compression_toolkit/ptq/runner.py,sha256=_c1dSjlPPpsx59Vbg1buhG9bZq__OORz1VlPkwjJzoc,2552
 model_compression_toolkit/ptq/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
@@ -507,7 +508,7 @@ model_compression_toolkit/xquant/common/core_report_generator.py,sha256=GHnJJpK6
 model_compression_toolkit/xquant/common/dataset_utils.py,sha256=91uXF9UwxdY7BvUT0FNkFm8a69c8oK8Xdl-y7lbuJxk,1649
 model_compression_toolkit/xquant/common/framework_report_utils.py,sha256=YE49232ESflW6ZaUABF1pk_GGHBxa_F1X5oRN2Jogys,3734
 model_compression_toolkit/xquant/common/model_analyzer.py,sha256=T_8OetIQNqR0nkfSatWsEceXSPYpHfYjboBPIyR03-w,3953
-model_compression_toolkit/xquant/common/model_folding_utils.py,sha256=y5Vmc-hJ2rJhzWdM53HdY-PrT5LlspejTUNlXaCrq9Q,4720
+model_compression_toolkit/xquant/common/model_folding_utils.py,sha256=7XMNmsngJgCPVjsuMNt6g4hzhkviB45qUmNRe9jQE7g,4815
 model_compression_toolkit/xquant/common/similarity_calculator.py,sha256=yCs_vlOThLzq7z-u2PkcEErLj7N7qCBPpRa6_5h34J8,10460
 model_compression_toolkit/xquant/common/similarity_functions.py,sha256=Atah1otdX9oUUch2JK-p-e291QHtkP_c4DfLG9WWo1Y,2935
 model_compression_toolkit/xquant/common/tensorboard_utils.py,sha256=85ABGQGKPZzctyZCHLazK0GxZ2ZUtQA3hZ_9fPiuMs0,6533
@@ -526,8 +527,8 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=bOc-hFL3gdoSM1Th_S2N_-9JJSlPGpZCTx_QLJHS6lg,3388
 model_compression_toolkit/xquant/pytorch/similarity_functions.py,sha256=CERxq5K8rqaiE-DlwhZBTUd9x69dtYJlkHOPLB54vm8,2354
 model_compression_toolkit/xquant/pytorch/tensorboard_utils.py,sha256=yjghWXxqOtT-QXoXBOuJyh45yUpFI0pKjdDegum2i68,9705
-mct_nightly-2.1.0.20240811.503.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.1.0.20240811.503.dist-info/METADATA,sha256=81BoRwx_nWGmZpRxLoEtAEKfBgJAIqvJ5ReNvEMB6Xg,19718
-mct_nightly-2.1.0.20240811.503.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-mct_nightly-2.1.0.20240811.503.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.1.0.20240811.503.dist-info/RECORD,,
+mct_nightly-2.1.0.20240812.432.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.1.0.20240812.432.dist-info/METADATA,sha256=2qqUeeA_e60PJ8S7cITQK2UlEOxMvWKKncdTK-H3v9E,19718
+mct_nightly-2.1.0.20240812.432.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+mct_nightly-2.1.0.20240812.432.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.1.0.20240812.432.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.1.0.20240811.000503"
+__version__ = "2.1.0.20240812.000432"

model_compression_toolkit/core/__init__.py

@@ -19,6 +19,7 @@ from model_compression_toolkit.core.common.quantization.debug_config import Debu
 from model_compression_toolkit.core.common.quantization import quantization_config
 from model_compression_toolkit.core.common.mixed_precision import mixed_precision_quantization_config
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig, QuantizationErrorMethod, DEFAULTCONFIG
+from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 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

model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_data.py

@@ -67,6 +67,7 @@ def compute_resource_utilization_data(in_model: Any,
                                                      fw_info,
                                                      fw_impl,
                                                      tpc,
+                                                     bit_width_config=core_config.bit_width_config,
                                                      mixed_precision_enable=mixed_precision_enable)
 
     # Compute parameters sum
@@ -227,6 +228,7 @@ def requires_mixed_precision(in_model: Any,
                                                  fw_info,
                                                  fw_impl,
                                                  tpc,
+                                                 bit_width_config=core_config.bit_width_config,
                                                  mixed_precision_enable=False)
     # Compute max weights memory in bytes
     weights_memory_by_layer_bytes, _ = compute_nodes_weights_params(transformed_graph, fw_info)

model_compression_toolkit/core/common/quantization/bit_width_config.py

@@ -0,0 +1,91 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from typing import List, Union, Dict
+
+from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.core.common.matchers.node_matcher import BaseNodeMatcher
+from model_compression_toolkit.logger import Logger
+
+
+class ManualBitWidthSelection:
+    """
+   Class to encapsulate the manual bit width selection configuration for a specific filter.
+
+   Attributes:
+       filter (BaseNodeMatcher): The filter used to select nodes for bit width manipulation.
+       bit_width (int): The bit width to be applied to the selected nodes.
+   """
+    def __init__(self,
+                 filter: BaseNodeMatcher,
+                 bit_width: int):
+        self.filter = filter
+        self.bit_width = bit_width
+
+
+class BitWidthConfig:
+    """
+    Class to manage manual bit-width configurations.
+
+    Attributes:
+        manual_activation_bit_width_selection_list (List[ManualBitWidthSelection]): A list of ManualBitWidthSelection objects defining manual bit-width configurations.
+    """
+    def __init__(self,
+                 manual_activation_bit_width_selection_list: List[ManualBitWidthSelection] = None):
+        self.manual_activation_bit_width_selection_list = [] if manual_activation_bit_width_selection_list is None else manual_activation_bit_width_selection_list
+
+    def __repr__(self):
+        # Used for debugging, thus no cover.
+        return str(self.__dict__)  # pragma: no cover
+
+    def set_manual_activation_bit_width(self,
+                                        filters: Union[List[BaseNodeMatcher], BaseNodeMatcher],
+                                        bit_widths: Union[List[int], int]):
+        """
+        Add a manual bit-width selection to the configuration.
+
+        Args:
+            filter (Union[List[BaseNodeMatcher], BaseNodeMatcher]): The filters used to select nodes for bit-width manipulation.
+            bit_width (Union[List[int], int]): The bit widths to be applied to the selected nodes.
+            If a single value is given it will be applied to all the filters
+        """
+        filters = [filters] if not isinstance(filters, list) else filters
+        bit_widths = [bit_widths] if not isinstance(bit_widths, list) else bit_widths
+        if len(bit_widths) > 1 and len(bit_widths) != len(filters):
+            Logger.critical(f"Configuration Error: The number of provided bit_width values {len(bit_widths)} "
+                            f"must match the number of filters {len(filters)}, or a single bit_width value "
+                            f"should be provided for all filters.")
+        elif len(bit_widths) == 1 and len(filters) > 1:
+            bit_widths = [bit_widths[0] for f in filters]
+        for bit_width, filter in zip (bit_widths, filters):
+            self.manual_activation_bit_width_selection_list += [ManualBitWidthSelection(filter, bit_width)]
+
+    def get_nodes_to_manipulate_bit_widths(self, graph: Graph) -> Dict:
+        """
+        Retrieve nodes from the graph that need their bit-widths changed according to the manual bit-width selections.
+
+        Args:
+            graph (Graph): The graph containing the nodes to be filtered and manipulated.
+
+        Returns:
+            Dict: A dictionary mapping nodes to their new bit-widths.
+        """
+        nodes_to_change_bit_width = {}
+        for manual_bit_width_selection in self.manual_activation_bit_width_selection_list:
+            filtered_nodes = graph.filter(manual_bit_width_selection.filter)
+            if len(filtered_nodes) == 0:
+                Logger.critical(f"Node Filtering Error: No nodes found in the graph for filter {manual_bit_width_selection.filter.__dict__} "
+                                f"to change their bit width to {manual_bit_width_selection.bit_width}.")
+            nodes_to_change_bit_width.update({n: manual_bit_width_selection.bit_width for n in filtered_nodes})
+        return nodes_to_change_bit_width

model_compression_toolkit/core/common/quantization/core_config.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
 from model_compression_toolkit.core.common.quantization.debug_config import DebugConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
@@ -22,9 +23,10 @@ class CoreConfig:
     A class to hold the configurations classes of the MCT-core.
     """
     def __init__(self,
-                 quantization_config: QuantizationConfig = QuantizationConfig(),
+                 quantization_config: QuantizationConfig = None,
                  mixed_precision_config: MixedPrecisionQuantizationConfig = None,
-                 debug_config: DebugConfig = DebugConfig()
+                 bit_width_config: BitWidthConfig = None,
+                 debug_config: DebugConfig = None
                  ):
         """
 
@@ -32,10 +34,12 @@ class CoreConfig:
             quantization_config (QuantizationConfig): Config for quantization.
             mixed_precision_config (MixedPrecisionQuantizationConfig): Config for mixed precision quantization.
             If None, a default MixedPrecisionQuantizationConfig is used.
+            bit_width_config (BitWidthConfig): Config for manual bit-width selection.
             debug_config (DebugConfig): Config for debugging and editing the network quantization process.
         """
-        self.quantization_config = quantization_config
-        self.debug_config = debug_config
+        self.quantization_config = QuantizationConfig() if quantization_config is None else quantization_config
+        self.bit_width_config = BitWidthConfig() if bit_width_config is None else bit_width_config
+        self.debug_config = DebugConfig() if debug_config is None else debug_config
 
         if mixed_precision_config is None:
             self.mixed_precision_config = MixedPrecisionQuantizationConfig()

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/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/pruning/keras/pruning_facade.py

@@ -21,6 +21,7 @@ from model_compression_toolkit.core.common.mixed_precision.resource_utilization_
 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/pruning/pytorch/pruning_facade.py

@@ -20,6 +20,7 @@ from model_compression_toolkit.core.common.mixed_precision.resource_utilization_
 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/xquant/common/model_folding_utils.py

@@ -16,6 +16,7 @@ from model_compression_toolkit.core.common.framework_implementation import Frame
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
+from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG
 
 from model_compression_toolkit.core.graph_prep_runner import graph_preparation_runner