mct-nightly 2.4.0.20250925.543__py3-none-any.whl → 2.4.2.20250926.532__py3-none-any.whl

model_compression_toolkit/gptq/runner.py

@@ -37,6 +37,7 @@ def _apply_gptq(gptq_config: GradientPTQConfig,
                 tb_w: TensorboardWriter,
                 tg: Graph,
                 tg_bias: Graph,
+                fw_info: FrameworkInfo,
                 fw_impl: FrameworkImplementation,
                 hessian_info_service: HessianInfoService = None) -> Graph:
     """
@@ -51,6 +52,7 @@ def _apply_gptq(gptq_config: GradientPTQConfig,
         tb_w: TensorBoardWriter object to log events.
         tg: Float Reference Graph.
         tg_bias: Graph of quantized model.
+        fw_info: 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: Framework implementation per framework
         hessian_info_service: HessianInfoService to fetch information based on the hessian approximation for the float model.
     Returns:
@@ -62,6 +64,7 @@ def _apply_gptq(gptq_config: GradientPTQConfig,
                                 gptq_config,
                                 representative_data_gen,
                                 fw_impl,
+                                fw_info,
                                 hessian_info_service=hessian_info_service)
 
         if tb_w is not None:
@@ -74,6 +77,7 @@ def gptq_runner(tg: Graph,
                 gptq_config: GradientPTQConfig,
                 representative_data_gen: Callable,
                 gptq_representative_data_gen: Callable,
+                fw_info: FrameworkInfo,
                 fw_impl: FrameworkImplementation,
                 tb_w: TensorboardWriter,
                 hessian_info_service: HessianInfoService = None) -> Graph:
@@ -87,6 +91,7 @@ def gptq_runner(tg: Graph,
         gptq_config: GradientPTQConfig with parameters about the tuning process.
         representative_data_gen: Dataset used for calibration.
         gptq_representative_data_gen: Dataset used for GPTQ training
+        fw_info: 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.
         tb_w: A TensorBoardWriter object initialized with the logger dir path if it was set, or None otherwise.
         hessian_info_service: HessianScoresService to fetch approximations of the hessian scores for the float model.
@@ -99,7 +104,7 @@ def gptq_runner(tg: Graph,
     #############################################
     # Apply Statistics Correction
     #############################################
-    tg_bias = apply_statistics_correction(tg, representative_data_gen, core_config, fw_impl, tb_w)
+    tg_bias = apply_statistics_correction(tg, representative_data_gen, core_config, fw_info, fw_impl, tb_w)
 
     if tb_w is not None:
         tb_w.add_graph(tg_bias, 'after_bias_correction')
@@ -112,6 +117,7 @@ def gptq_runner(tg: Graph,
                           tb_w,
                           tg,
                           tg_bias,
+                          fw_info,
                           fw_impl,
                           hessian_info_service=hessian_info_service)
 

model_compression_toolkit/pruning/keras/pruning_facade.py

@@ -17,7 +17,6 @@ from typing import Callable, Tuple, Union
 
 from model_compression_toolkit import get_target_platform_capabilities
 from model_compression_toolkit.constants import TENSORFLOW
-from model_compression_toolkit.quantization_preparation.load_fqc import load_fqc_configuration
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
 from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
@@ -25,8 +24,10 @@ 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.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
+from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG
 from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
 
 if FOUND_TF:
@@ -34,12 +35,11 @@ if FOUND_TF:
         AttachTpcToKeras
     from model_compression_toolkit.core.keras.back2framework.float_model_builder import FloatKerasModelBuilder
     from model_compression_toolkit.core.keras.pruning.pruning_keras_implementation import PruningKerasImplementation
-    from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
+    from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from tensorflow.keras.models import Model
 
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
-    @set_keras_info
     def keras_pruning_experimental(model: Model,
                                    target_resource_utilization: ResourceUtilization,
                                    representative_data_gen: Callable,
@@ -116,25 +116,30 @@ if FOUND_TF:
 
         target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
-        framework_platform_capabilities = AttachTpcToKeras().attach(target_platform_capabilities)
+        attach2keras = AttachTpcToKeras()
+        target_platform_capabilities = attach2keras.attach(target_platform_capabilities)
 
         # Convert the original Keras model to an internal graph representation.
         float_graph = read_model_to_graph(model,
                                           representative_data_gen,
-                                          framework_platform_capabilities,
+                                          target_platform_capabilities,
+                                          DEFAULT_KERAS_INFO,
                                           fw_impl)
 
         # Apply quantization configuration to the graph. This step is necessary even when not quantizing,
         # as it prepares the graph for the pruning process.
-        float_graph_with_compression_config = load_fqc_configuration(float_graph, framework_platform_capabilities)
+        float_graph_with_compression_config = set_quantization_configuration_to_graph(float_graph,
+                                                                                      quant_config=DEFAULTCONFIG,
+                                                                                      mixed_precision_enable=False)
 
         # Create a Pruner object with the graph and configuration.
         pruner = Pruner(float_graph_with_compression_config,
+                        DEFAULT_KERAS_INFO,
                         fw_impl,
                         target_resource_utilization,
                         representative_data_gen,
                         pruning_config,
-                        framework_platform_capabilities)
+                        target_platform_capabilities)
 
         # Apply the pruning process.
         pruned_graph = pruner.prune_graph()

model_compression_toolkit/pruning/pytorch/pruning_facade.py

@@ -23,9 +23,10 @@ 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.quantization_preparation.load_fqc import load_fqc_configuration
+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
+from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG
 from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
 
 
@@ -35,7 +36,7 @@ if FOUND_TORCH:
     from model_compression_toolkit.core.pytorch.back2framework.float_model_builder import FloatPyTorchModelBuilder
     from model_compression_toolkit.core.pytorch.pruning.pruning_pytorch_implementation import \
         PruningPytorchImplementation
-    from model_compression_toolkit.core.pytorch.default_framework_info import set_pytorch_info
+    from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
     from torch.nn import Module
     from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2pytorch import \
         AttachTpcToPytorch
@@ -43,7 +44,6 @@ if FOUND_TORCH:
     # Set the default Target Platform Capabilities (TPC) for PyTorch.
     DEFAULT_PYOTRCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
 
-    @set_pytorch_info
     def pytorch_pruning_experimental(model: Module,
                                      target_resource_utilization: ResourceUtilization,
                                      representative_data_gen: Callable,
@@ -129,14 +129,18 @@ if FOUND_TORCH:
         float_graph = read_model_to_graph(model,
                                           representative_data_gen,
                                           framework_platform_capabilities,
+                                          DEFAULT_PYTORCH_INFO,
                                           fw_impl)
 
         # Apply quantization configuration to the graph. This step is necessary even when not quantizing,
         # as it prepares the graph for the pruning process.
-        float_graph_with_compression_config = load_fqc_configuration(float_graph, framework_platform_capabilities)
+        float_graph_with_compression_config = set_quantization_configuration_to_graph(float_graph,
+                                                                                      quant_config=DEFAULTCONFIG,
+                                                                                      mixed_precision_enable=False)
 
         # Create a Pruner object with the graph and configuration.
         pruner = Pruner(float_graph_with_compression_config,
+                        DEFAULT_PYTORCH_INFO,
                         fw_impl,
                         target_resource_utilization,
                         representative_data_gen,

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -36,8 +36,9 @@ from model_compression_toolkit.metadata import create_model_metadata
 if FOUND_TF:
     from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
         AttachTpcToKeras
-    from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
+    from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
+    from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.exporter.model_wrapper import get_exportable_keras_model
@@ -48,7 +49,6 @@ if FOUND_TF:
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
 
-    @set_keras_info
     def keras_post_training_quantization(in_model: Model,
                                          representative_data_gen: Callable,
                                          target_resource_utilization: ResourceUtilization = None,
@@ -121,20 +121,25 @@ if FOUND_TF:
 
             >>> quantized_model, quantization_info = mct.ptq.keras_post_training_quantization(model, repr_datagen, ru, core_config=config)
 
-            For more configuration options, please take a look at our `API documentation <https://sonysemiconductorsolutions.github.io/mct-model-optimization/api/api_docs/modules/mixed_precision_quantization_config.html>`_.
+            For more configuration options, please take a look at our `API documentation <https://sony.github.io/model_optimization/api/api_docs/modules/mixed_precision_quantization_config.html>`_.
 
          """
 
         if core_config.debug_config.bypass:
             return in_model, None
 
+        fw_info = DEFAULT_KERAS_INFO
+
+        KerasModelValidation(model=in_model,
+                             fw_info=fw_info).validate()
+
         if core_config.is_mixed_precision_enabled:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
                 Logger.critical("Given quantization config to mixed-precision facade is not of type "
                                 "MixedPrecisionQuantizationConfig. Please use keras_post_training_quantization "
                                 "API, or pass a valid mixed precision configuration.")  # pragma: no cover
 
-        tb_w = init_tensorboard_writer()
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = KerasImplementation()
 
@@ -148,6 +153,7 @@ if FOUND_TF:
         tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_model,
                                                                 representative_data_gen=representative_data_gen,
                                                                 core_config=core_config,
+                                                                fw_info=fw_info,
                                                                 fw_impl=fw_impl,
                                                                 fqc=framework_platform_capabilities,
                                                                 target_resource_utilization=target_resource_utilization,
@@ -163,6 +169,7 @@ if FOUND_TF:
         graph_with_stats_correction = ptq_runner(tg,
                                                  representative_data_gen,
                                                  core_config,
+                                                 fw_info,
                                                  fw_impl,
                                                  tb_w)
 
@@ -172,7 +179,8 @@ if FOUND_TF:
                                         tb_w,
                                         similarity_baseline_graph,
                                         quantized_graph,
-                                        fw_impl)
+                                        fw_impl,
+                                        fw_info)
 
         exportable_model, user_info = get_exportable_keras_model(graph_with_stats_correction)
         if framework_platform_capabilities.tpc.add_metadata:

model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -34,7 +34,7 @@ from model_compression_toolkit.core.common.quantization.quantize_graph_weights i
 from model_compression_toolkit.metadata import create_model_metadata
 
 if FOUND_TORCH:
-    from model_compression_toolkit.core.pytorch.default_framework_info import set_pytorch_info
+    from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
     from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from torch.nn import Module
@@ -46,7 +46,6 @@ if FOUND_TORCH:
 
     DEFAULT_PYTORCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
 
-    @set_pytorch_info
     def pytorch_post_training_quantization(in_module: Module,
                                            representative_data_gen: Callable,
                                            target_resource_utilization: ResourceUtilization = None,
@@ -103,6 +102,8 @@ if FOUND_TORCH:
         if core_config.debug_config.bypass:
             return in_module, None
 
+        fw_info = DEFAULT_PYTORCH_INFO
+
         if core_config.is_mixed_precision_enabled:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
                 Logger.critical("Given quantization config to mixed-precision facade is not of type "
@@ -110,7 +111,7 @@ if FOUND_TORCH:
                                 "pytorch_post_training_quantization API, or pass a valid mixed precision "
                                 "configuration.")  # pragma: no cover
 
-        tb_w = init_tensorboard_writer()
+        tb_w = init_tensorboard_writer(fw_info)
 
         fw_impl = PytorchImplementation()
 
@@ -124,6 +125,7 @@ if FOUND_TORCH:
         tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_module,
                                                                 representative_data_gen=representative_data_gen,
                                                                 core_config=core_config,
+                                                                fw_info=fw_info,
                                                                 fw_impl=fw_impl,
                                                                 fqc=framework_platform_capabilities,
                                                                 target_resource_utilization=target_resource_utilization,
@@ -139,6 +141,7 @@ if FOUND_TORCH:
         graph_with_stats_correction = ptq_runner(tg,
                                                  representative_data_gen,
                                                  core_config,
+                                                 fw_info,
                                                  fw_impl,
                                                  tb_w)
 
@@ -148,7 +151,8 @@ if FOUND_TORCH:
                                         tb_w,
                                         similarity_baseline_graph,
                                         quantized_graph,
-                                        fw_impl)
+                                        fw_impl,
+                                        fw_info)
 
         exportable_model, user_info = get_exportable_pytorch_model(graph_with_stats_correction)
         if framework_platform_capabilities.tpc.add_metadata:

model_compression_toolkit/ptq/runner.py

@@ -16,6 +16,7 @@
 
 from typing import Callable
 
+from model_compression_toolkit.core.common import FrameworkInfo
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
@@ -27,6 +28,7 @@ from model_compression_toolkit.core.common.visualization.tensorboard_writer impo
 def ptq_runner(tg: Graph,
                representative_data_gen: Callable,
                core_config: CoreConfig,
+               fw_info: FrameworkInfo,
                fw_impl: FrameworkImplementation,
                tb_w: TensorboardWriter) -> Graph:
     """
@@ -36,6 +38,7 @@ def ptq_runner(tg: Graph,
         tg: Graph to apply PTQ and to quantize.
         representative_data_gen (Callable): Dataset used for calibration.
         core_config: CoreConfig containing parameters of how the model should be quantized.
+        fw_info: 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.
         tb_w: A TensorBoardWriter object initialized with the logger dir path if it was set, or None otherwise.
@@ -47,5 +50,5 @@ def ptq_runner(tg: Graph,
     #############################################
     # Statistics Correction
     #############################################
-    tg = apply_statistics_correction(tg, representative_data_gen, core_config, fw_impl, tb_w)
+    tg = apply_statistics_correction(tg, representative_data_gen, core_config, fw_info, fw_impl, tb_w)
     return tg

model_compression_toolkit/qat/common/qat_config.py

@@ -19,17 +19,21 @@ from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.trainable_infrastructure import TrainingMethod
 
 
-def is_qat_applicable(node: common.BaseNode) -> bool:
+def is_qat_applicable(node: common.BaseNode,
+                      fw_info: FrameworkInfo) -> bool:
     """
     A function for deciding if a layer should be fine-tuned during QAT
 
     Args:
         node (BaseNode): Node for quantization decision
+        fw_info (FrameworkInfo): Pytorch quantization information
 
     Returns:
         A boolean whether the layer is to be wrapped with a QuantizeWrapper
     """
-    return (node.kernel_attr is not None and node.is_weights_quantization_enabled(node.kernel_attr)) \
+
+    kernel_attr = fw_info.get_kernel_op_attributes(node.type)[0]
+    return (kernel_attr is not None and node.is_weights_quantization_enabled(kernel_attr)) \
             or node.is_activation_quantization_enabled()
 
 

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -37,9 +37,10 @@ if FOUND_TF:
     from tensorflow.keras.models import Model
 
     from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
+    from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
+    from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
-    from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
 
     from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
 
@@ -51,6 +52,7 @@ if FOUND_TF:
     from model_compression_toolkit.constants import TENSORFLOW
     from model_compression_toolkit.qat.common.qat_config import is_qat_applicable
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
+    from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.qat.keras.quantizer.quantization_builder import quantization_builder, \
     get_activation_quantizer_holder
     from model_compression_toolkit.qat.common.qat_config import QATConfig
@@ -71,11 +73,11 @@ if FOUND_TF:
         Returns: Wrapped layer
 
         """
-        if is_qat_applicable(n):
+        if is_qat_applicable(n, DEFAULT_KERAS_INFO):
             # If we are here, then the node has a kernel attribute to quantize and training during QAT
             weights_quantizers, _ = quantization_builder(n,
                                                          qat_config,
-                                                         n.kernel_attr)
+                                                         DEFAULT_KERAS_INFO.get_kernel_op_attributes(n.type)[0])
             if len(weights_quantizers) > 0:
                 layer.trainable = True
                 return KerasTrainableQuantizationWrapper(layer, weights_quantizers)
@@ -85,7 +87,6 @@ if FOUND_TF:
         return layer
 
 
-    @set_keras_info
     def keras_quantization_aware_training_init_experimental(in_model: Model,
                                                             representative_data_gen: Callable,
                                                             target_resource_utilization: ResourceUtilization = None,
@@ -166,7 +167,7 @@ if FOUND_TF:
 
              >>> quantized_model = tf.keras.models.load_model(model_file, custom_objects=custom_objects)
 
-             For more configuration options, please take a look at our `API documentation <https://sonysemiconductorsolutions.github.io/mct-model-optimization/api/api_docs/modules/mixed_precision_quantization_config.html>`_.
+             For more configuration options, please take a look at our `API documentation <https://sony.github.io/model_optimization/api/api_docs/modules/mixed_precision_quantization_config.html>`_.
 
          """
 
@@ -174,13 +175,16 @@ if FOUND_TF:
                        f"If you encounter an issue, please open an issue in our GitHub "
                        f"project https://github.com/sony/model_optimization")
 
+        KerasModelValidation(model=in_model,
+                             fw_info=DEFAULT_KERAS_INFO).validate()
+
         if core_config.is_mixed_precision_enabled:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
                 Logger.critical("Given quantization config to mixed-precision facade is not of type "
                              "MixedPrecisionQuantizationConfig. Please use keras_post_training_quantization API,"
                              "or pass a valid mixed precision configuration.")
 
-        tb_w = init_tensorboard_writer()
+        tb_w = init_tensorboard_writer(DEFAULT_KERAS_INFO)
 
         fw_impl = KerasImplementation()
 
@@ -194,15 +198,17 @@ if FOUND_TF:
         tg, bit_widths_config, _, _ = core_runner(in_model=in_model,
                                                   representative_data_gen=representative_data_gen,
                                                   core_config=core_config,
+                                                  fw_info=DEFAULT_KERAS_INFO,
                                                   fw_impl=fw_impl,
                                                   fqc=target_platform_capabilities,
                                                   target_resource_utilization=target_resource_utilization,
                                                   tb_w=tb_w)
 
-        tg = ptq_runner(tg, representative_data_gen, core_config, fw_impl, tb_w)
+        tg = ptq_runner(tg, representative_data_gen, core_config, DEFAULT_KERAS_INFO, fw_impl, tb_w)
 
         _qat_wrapper = partial(qat_wrapper, qat_config=qat_config)
         qat_model, user_info = KerasModelBuilder(graph=tg,
+                                                 fw_info=DEFAULT_KERAS_INFO,
                                                  wrapper=_qat_wrapper,
                                                  get_activation_quantizer_holder_fn=partial(get_activation_quantizer_holder,
                                                                                             qat_config=qat_config)).build_model()

model_compression_toolkit/qat/pytorch/quantization_facade.py

@@ -36,7 +36,7 @@ if FOUND_TORCH:
     import torch.nn as nn
     from torch.nn import Module
     from mct_quantizers import PytorchActivationQuantizationHolder
-    from model_compression_toolkit.core.pytorch.default_framework_info import set_pytorch_info
+    from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
     from model_compression_toolkit.qat.common.qat_config import is_qat_applicable
@@ -62,10 +62,10 @@ if FOUND_TORCH:
         Returns: Wrapped layer
 
         """
-        if is_qat_applicable(n):
+        if is_qat_applicable(n, DEFAULT_PYTORCH_INFO):
             # If we are here, then the node has a kernel attribute to quantize and training during QAT
             weights_quantizers, _ = quantization_builder(n, qat_config,
-                                                         n.kernel_attr)
+                                                         DEFAULT_PYTORCH_INFO.get_kernel_op_attributes(n.type)[0])
             if len(weights_quantizers) > 0:
                 return PytorchQuantizationWrapper(module, weights_quantizers)
 
@@ -74,7 +74,6 @@ if FOUND_TORCH:
         return module
 
 
-    @set_pytorch_info
     def pytorch_quantization_aware_training_init_experimental(in_model: Module,
                                                               representative_data_gen: Callable,
                                                               target_resource_utilization: ResourceUtilization = None,
@@ -136,7 +135,7 @@ if FOUND_TORCH:
 
              >>> quantized_model, quantization_info = mct.qat.pytorch_quantization_aware_training_init_experimental(model, repr_datagen, core_config=config)
 
-             For more configuration options, please take a look at our `API documentation <https://sonysemiconductorsolutions.github.io/mct-model-optimization/api/api_docs/modules/mixed_precision_quantization_config.html>`_.
+             For more configuration options, please take a look at our `API documentation <https://sony.github.io/model_optimization/api/api_docs/modules/mixed_precision_quantization_config.html>`_.
 
          """
         Logger.warning(
@@ -150,7 +149,7 @@ if FOUND_TORCH:
                                 "MixedPrecisionQuantizationConfig. Please use pytorch_post_training_quantization API,"
                                 "or pass a valid mixed precision configuration.")
 
-        tb_w = init_tensorboard_writer()
+        tb_w = init_tensorboard_writer(DEFAULT_PYTORCH_INFO)
         fw_impl = PytorchImplementation()
 
         target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
@@ -163,16 +162,18 @@ if FOUND_TORCH:
         tg, bit_widths_config, _, _ = core_runner(in_model=in_model,
                                                   representative_data_gen=representative_data_gen,
                                                   core_config=core_config,
+                                                  fw_info=DEFAULT_PYTORCH_INFO,
                                                   fw_impl=fw_impl,
                                                   fqc=framework_platform_capabilities,
                                                   target_resource_utilization=target_resource_utilization,
                                                   tb_w=tb_w)
 
-        tg = ptq_runner(tg, representative_data_gen, core_config, fw_impl, tb_w)
+        tg = ptq_runner(tg, representative_data_gen, core_config, DEFAULT_PYTORCH_INFO, fw_impl, tb_w)
 
         _qat_wrapper = partial(qat_wrapper, qat_config=qat_config)
 
         qat_model, user_info = PyTorchModelBuilder(graph=tg,
+                                                   fw_info=DEFAULT_PYTORCH_INFO,
                                                    wrapper=_qat_wrapper,
                                                    get_activation_quantizer_holder_fn=partial(
                                                        get_activation_quantizer_holder,
@@ -180,6 +181,9 @@ if FOUND_TORCH:
 
         user_info.mixed_precision_cfg = bit_widths_config
 
+        # Remove fw_info from graph to enable saving the pytorch model (fw_info can not be pickled)
+        delattr(qat_model.graph, 'fw_info')
+
         return qat_model, user_info
 
 

model_compression_toolkit/target_platform_capabilities/constants.py

@@ -29,7 +29,7 @@ QNNPACK_TP_MODEL = 'qnnpack'
 # TP Attributes
 KERNEL_ATTR = "kernel_attr"
 BIAS_ATTR = "bias_attr"
-POSITIONAL_ATTR = "pos_attr"
+POS_ATTR = "pos_attr"
 
 # TODO: this is duplicated from the core frameworks constants files, because the original consts can't be used here
 #  duo to circular dependency. It might be best to extract the constants from the core file and put them here (in a

model_compression_toolkit/target_platform_capabilities/targetplatform2framework/attach2pytorch.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+# Copyright 2025 Sony Semiconductor Solutions, 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.
@@ -32,7 +32,7 @@ from model_compression_toolkit.target_platform_capabilities.targetplatform2frame
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2fw import \
     AttachTpcToFramework
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attribute_filter import Eq
-from edgemdt_cl.pytorch import MulticlassNMS, MulticlassNMSWithIndices
+from edgemdt_cl.pytorch import MulticlassNMS, MulticlassNMSWithIndices, MulticlassNMSOBB
 
 
 class AttachTpcToPytorch(AttachTpcToFramework):
@@ -98,7 +98,7 @@ class AttachTpcToPytorch(AttachTpcToFramework):
             OperatorSetNames.L2NORM: [LayerFilterParams(torch.nn.functional.normalize,
                                                         Eq('p', 2) | Eq('p', None))],
             OperatorSetNames.SSD_POST_PROCESS: [],  # no such operator in pytorch
-            OperatorSetNames.COMBINED_NON_MAX_SUPPRESSION: [MulticlassNMS, MulticlassNMSWithIndices],
+            OperatorSetNames.COMBINED_NON_MAX_SUPPRESSION: [MulticlassNMS, MulticlassNMSWithIndices, MulticlassNMSOBB],
             OperatorSetNames.EXP: [torch.exp],
             OperatorSetNames.SIN: [torch.sin],
             OperatorSetNames.COS: [torch.cos],

model_compression_toolkit/trainable_infrastructure/common/get_quantizer_config.py

@@ -48,6 +48,7 @@ def get_trainable_quantizer_weights_config(
                                            final_attr_cfg.enable_weights_quantization,
                                            final_attr_cfg.weights_channels_axis[0],  # Output channel axis
                                            final_attr_cfg.weights_per_channel_threshold,
+                                           final_node_cfg.min_threshold,
                                            weights_quantization_candidates)
 
 
@@ -75,6 +76,7 @@ def get_trainable_quantizer_activation_config(
                                               final_cfg.activation_n_bits,
                                               final_cfg.activation_quantization_params,
                                               final_cfg.enable_activation_quantization,
+                                              final_cfg.min_threshold,
                                               activation_quantization_candidates)
 
 

model_compression_toolkit/trainable_infrastructure/common/trainable_quantizer_config.py

@@ -44,6 +44,7 @@ class TrainableQuantizerActivationConfig:
                  activation_n_bits: int,
                  activation_quantization_params: Dict,
                  enable_activation_quantization: bool,
+                 min_threshold: float,
                  activation_quantization_candidates: List[TrainableQuantizerCandidateConfig] = None,
                  ):
         """
@@ -54,11 +55,13 @@ class TrainableQuantizerActivationConfig:
             activation_n_bits (int): Number of bits to quantize the activations.
             activation_quantization_params (Dict): Dictionary that contains activation quantization params.
             enable_activation_quantization (bool): Whether to quantize the layer's activations or not.
+            min_threshold (float): Minimum threshold to use during thresholds selection.
         """
         self.activation_quantization_method = activation_quantization_method
         self.activation_n_bits = activation_n_bits
         self.activation_quantization_params = activation_quantization_params
         self.enable_activation_quantization = enable_activation_quantization
+        self.min_threshold = min_threshold
         self.activation_bits_candidates = activation_quantization_candidates
 
 
@@ -70,6 +73,7 @@ class TrainableQuantizerWeightsConfig:
                  enable_weights_quantization: bool,
                  weights_channels_axis: int,
                  weights_per_channel_threshold: bool,
+                 min_threshold: float,
                  weights_quantization_candidates: List[TrainableQuantizerCandidateConfig] = None,
                  ):
         """
@@ -82,6 +86,7 @@ class TrainableQuantizerWeightsConfig:
             enable_weights_quantization (bool): Whether to quantize the layer's weights or not.
             weights_channels_axis (int): Axis to quantize a node's kernel when quantizing per-channel.
             weights_per_channel_threshold (bool): Whether to quantize the weights per-channel or not (per-tensor).
+            min_threshold (float): Minimum threshold to use during thresholds selection.
         """
         self.weights_quantization_method = weights_quantization_method
         self.weights_n_bits = weights_n_bits
@@ -89,4 +94,5 @@ class TrainableQuantizerWeightsConfig:
         self.enable_weights_quantization = enable_weights_quantization
         self.weights_channels_axis = weights_channels_axis
         self.weights_per_channel_threshold = weights_per_channel_threshold
+        self.min_threshold = min_threshold
         self.weights_bits_candidates = weights_quantization_candidates

model_compression_toolkit/trainable_infrastructure/keras/config_serialization.py

@@ -77,11 +77,13 @@ def config_deserialization(in_config: dict) -> Union[TrainableQuantizerWeightsCo
                                                weights_quantization_params=weights_quantization_params,
                                                enable_weights_quantization=in_config[C.ENABLE_WEIGHTS_QUANTIZATION],
                                                weights_channels_axis=in_config[C.WEIGHTS_CHANNELS_AXIS],
-                                               weights_per_channel_threshold=in_config[C.WEIGHTS_PER_CHANNEL_THRESHOLD])
+                                               weights_per_channel_threshold=in_config[C.WEIGHTS_PER_CHANNEL_THRESHOLD],
+                                               min_threshold=in_config[C.MIN_THRESHOLD])
     elif in_config[C.IS_ACTIVATIONS]:
         return TrainableQuantizerActivationConfig(activation_quantization_method=QuantizationMethod(in_config[C.ACTIVATION_QUANTIZATION_METHOD]),
                                                   activation_n_bits=in_config[C.ACTIVATION_N_BITS],
                                                   activation_quantization_params=in_config[C.ACTIVATION_QUANTIZATION_PARAMS],
-                                                  enable_activation_quantization=in_config[C.ENABLE_ACTIVATION_QUANTIZATION])
+                                                  enable_activation_quantization=in_config[C.ENABLE_ACTIVATION_QUANTIZATION],
+                                                  min_threshold=in_config[C.MIN_THRESHOLD])
     else:
         raise NotImplemented  # pragma: no cover

model_compression_toolkit/xquant/__init__.py

@@ -16,4 +16,5 @@
 from model_compression_toolkit.xquant.common.xquant_config import XQuantConfig
 from model_compression_toolkit.xquant.keras.facade_xquant_report import xquant_report_keras_experimental
 from model_compression_toolkit.xquant.pytorch.facade_xquant_report import xquant_report_pytorch_experimental
+from model_compression_toolkit.xquant.pytorch.facade_xquant_report import xquant_report_troubleshoot_pytorch_experimental