mct-nightly 2.2.0.20240908.528__py3-none-any.whl → 2.2.0.20240909.453__py3-none-any.whl

{mct_nightly-2.2.0.20240908.528.dist-info → mct_nightly-2.2.0.20240909.453.dist-info}/METADATA

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

{mct_nightly-2.2.0.20240908.528.dist-info → mct_nightly-2.2.0.20240909.453.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=gl_OAgmU-izJxoKGJjAirXzkQIzQU8OSKEdHXUwe2to,1573
+model_compression_toolkit/__init__.py,sha256=msE8-6YmO3FQaiZRSwN96Pup9KiZoPuXlMDgqzD6fEI,1573
 model_compression_toolkit/constants.py,sha256=i4wYheBkIdQmsQA-axIpcT3YiSO1USNc-jaNiNE8w6E,3920
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -375,7 +375,7 @@ model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=TMus5LYJnTn
 model_compression_toolkit/gptq/pytorch/quantizer/__init__.py,sha256=ZHNHo1yzye44m9_ht4UUZfTpK01RiVR3Tr74-vtnOGI,968
 model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py,sha256=fKg-PNOhGBiL-4eySS9Fyw0GkA76Pq8jT_HbJuJ8iZU,4143
 model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py,sha256=OocYYRqvl7rZ37QT0hTzfJnWGiNCPskg7cziTlR7TRk,3893
-model_compression_toolkit/gptq/pytorch/quantizer/quantization_builder.py,sha256=Lf334209uVFXuRKIFqVvq9RyEcv014Bozt1hr_O6XjQ,4447
+model_compression_toolkit/gptq/pytorch/quantizer/quantization_builder.py,sha256=5EyAzvlU01vLyXmMwY_8dNyb7GwYktXmnrvUON8n8WI,4696
 model_compression_toolkit/gptq/pytorch/quantizer/regularization_factory.py,sha256=mDWZERLwtDzqWeJUwHMVyGdlS8wPLjJ3NvZiKBP6BNA,1959
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/__init__.py,sha256=lNJ29DYxaLUPDstRDA1PGI5r9Fulq_hvrZMlhst1Z5g,697
 model_compression_toolkit/gptq/pytorch/quantizer/soft_rounding/soft_quantizer_reg.py,sha256=oO7WgsAHMnWoXNm_gTKAAe-Nd79mGL_m677ai-ui424,4132
@@ -486,7 +486,7 @@ model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/
 model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tpc_pytorch.py,sha256=YVJJvqGPBdkKnug99p9bjqtbfecDXZKIB2iWVCe7RUY,5960
 model_compression_toolkit/trainable_infrastructure/__init__.py,sha256=uewpvlPkH9mBFt8IxoAgIfz6iEcvWbOImm_fb6_BxD8,1543
 model_compression_toolkit/trainable_infrastructure/common/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
-model_compression_toolkit/trainable_infrastructure/common/base_trainable_quantizer.py,sha256=i5ZX0UnSt_XAgxGyyd7ZRHcocuwTh_FxWgGD2qN7zFc,7735
+model_compression_toolkit/trainable_infrastructure/common/base_trainable_quantizer.py,sha256=IF50ASBUvVrOVqlJ1nHNxZxKXSuCanjhUX0YjMB-rRg,7946
 model_compression_toolkit/trainable_infrastructure/common/constants.py,sha256=HN120boJxAnEXNrLSj-o_s-VX4o6C-1ap_KZ4840sd0,875
 model_compression_toolkit/trainable_infrastructure/common/get_quantizer_config.py,sha256=Jxd4IjS_t0FwnA_S_WmZeVbh4VM6Da9ahKGPLp6ZhQo,6983
 model_compression_toolkit/trainable_infrastructure/common/get_quantizers.py,sha256=KoX-6LJMsRzXy0i72ve4buJ32cGNQVHVLqHJxhv0lPQ,3428
@@ -500,7 +500,7 @@ model_compression_toolkit/trainable_infrastructure/keras/load_model.py,sha256=DJ
 model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha256=eVB5FSE3OmTLrhfLUcP2knwN1z2_unQLM-xFEGwdafA,5587
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
-model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=7ZFf_E8nFao5f38Qk4-GzGxHgrKTHGj-4ohgPzq2Z7k,2304
+model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=lWc5EG3ptrP85n69EHGKFkIadnrKEBMKnB5YXQ5AmXo,2745
 model_compression_toolkit/trainable_infrastructure/pytorch/quantizer_utils.py,sha256=1yOXKghUYfw2hmzbqTuNagIXBoM-wR2bP-ul66-mnDw,7767
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/__init__.py,sha256=73CXhqqNTvDpsvlJXclrGJq-vsCUYCI64ILu1y2mtvw,1056
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/base_activation_quantizer.py,sha256=X6E6mewWQot_aAkz3UxW5X0-Fjl_aMMjs3A-Af5eL6w,972
@@ -508,8 +508,8 @@ model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/lsq/symmetric_lsq.py,sha256=0UGoFHAR-RP9aFbAOILbM8kAG9OwUJJZ_g3Rz58SGlY,5462
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/lsq/uniform_lsq.py,sha256=BPeunWrYNmbduZGXiZKy5t1ubYREX7QqWOXv2Dt85lk,5285
 model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/ste/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
-model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/ste/symmetric_ste.py,sha256=20DEZgn6ZepcjKrATvciaiQNs2VGf5uwF6f6hDJLOVo,5226
-model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/ste/uniform_ste.py,sha256=1XHClqM7EhNvYiH6sqs6OI3JUGPfjW55v2eQotVwy8c,5010
+model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/ste/symmetric_ste.py,sha256=p11HY91muyovTXaKLEPQ48WOi8ge_Z1K7KdX9Y56mgw,5443
+model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/ste/uniform_ste.py,sha256=8zjzP-dxF1FEQ-qnFcFhWpbndNeVtvaPL3uEPQpbdLk,5202
 model_compression_toolkit/xquant/__init__.py,sha256=vdmr8sQw3jIBLF9ck7qrskPoXzDKtksHWlMOkU1JUnQ,1003
 model_compression_toolkit/xquant/common/__init__.py,sha256=ycb1Xt7PtixY2Uabr94JGSwBMcct66O8ZMVf3Qa3ud8,719
 model_compression_toolkit/xquant/common/constants.py,sha256=k-9LOEv1n_m8dV4chX0dNOTWyhhF7S00E0lkUxtO84E,1592
@@ -536,8 +536,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=mkoEktLFFHtEKzzFRn_jCnxjhJolK12TZ5AQeDHzUO8,9767
-mct_nightly-2.2.0.20240908.528.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20240908.528.dist-info/METADATA,sha256=BOMyD_uhRmYOemgFCIH970uj81mEPJCLr_akH0WuzcY,20813
-mct_nightly-2.2.0.20240908.528.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-mct_nightly-2.2.0.20240908.528.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.2.0.20240908.528.dist-info/RECORD,,
+mct_nightly-2.2.0.20240909.453.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20240909.453.dist-info/METADATA,sha256=smPztOdlwS0FQKktb9hXWTvuuliw7uO4689U-ondGtg,20813
+mct_nightly-2.2.0.20240909.453.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+mct_nightly-2.2.0.20240909.453.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.2.0.20240909.453.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.2.0.20240908.000528"
+__version__ = "2.2.0.20240909.000453"

model_compression_toolkit/gptq/pytorch/quantizer/quantization_builder.py

@@ -25,8 +25,9 @@ from mct_quantizers.common.get_quantizers import get_inferable_quantizer_class
 from mct_quantizers.pytorch.quantizers import BasePyTorchInferableQuantizer
 
 from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.trainable_infrastructure import TrainingMethod, BasePytorchActivationTrainableQuantizer
 from model_compression_toolkit.trainable_infrastructure.common.get_quantizer_config import \
-    get_trainable_quantizer_weights_config
+    get_trainable_quantizer_weights_config, get_trainable_quantizer_activation_config
 from model_compression_toolkit.trainable_infrastructure.common.get_quantizers import \
     get_trainable_quantizer_class
 
@@ -68,12 +69,11 @@ def quantization_builder(n: common.BaseNode,
 
         quant_method = n.final_activation_quantization_cfg.activation_quantization_method
 
-        quantizer_class = get_inferable_quantizer_class(quant_target=QuantizationTarget.Activation,
+        quantizer_class = get_trainable_quantizer_class(quant_target=QuantizationTarget.Activation,
+                                                        quantizer_id=TrainingMethod.STE,
                                                         quant_method=quant_method,
-                                                        quantizer_base_class=BasePyTorchInferableQuantizer)
-
-        kwargs = get_activation_inferable_quantizer_kwargs(n.final_activation_quantization_cfg)
-
-        activation_quantizers.append(quantizer_class(**kwargs))
+                                                        quantizer_base_class=BasePytorchActivationTrainableQuantizer)
+        cfg = get_trainable_quantizer_activation_config(n, None)
+        activation_quantizers.append(quantizer_class(cfg, freeze_quant_params=True))
 
     return weights_quantizers, activation_quantizers

model_compression_toolkit/trainable_infrastructure/common/base_trainable_quantizer.py

@@ -14,18 +14,16 @@
 # ==============================================================================
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Union, List, Any
 from inspect import signature
-
-from model_compression_toolkit.logger import Logger
+from typing import Union, List, Any
 
 from mct_quantizers.common.base_inferable_quantizer import BaseInferableQuantizer, \
     QuantizationTarget
-from model_compression_toolkit.trainable_infrastructure.common.trainable_quantizer_config import \
-    TrainableQuantizerActivationConfig, TrainableQuantizerWeightsConfig
 from mct_quantizers.common.constants import QUANTIZATION_METHOD, \
     QUANTIZATION_TARGET
-
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.trainable_infrastructure.common.trainable_quantizer_config import \
+    TrainableQuantizerActivationConfig, TrainableQuantizerWeightsConfig
 
 VAR = 'var'
 GROUP = 'group'
@@ -43,12 +41,14 @@ class VariableGroup(Enum):
 
 class BaseTrainableQuantizer(BaseInferableQuantizer, ABC):
     def __init__(self,
-                 quantization_config: Union[TrainableQuantizerActivationConfig, TrainableQuantizerWeightsConfig]):
+                 quantization_config: Union[TrainableQuantizerActivationConfig, TrainableQuantizerWeightsConfig],
+                 freeze_quant_params: bool = False):
         """
         This class is a base quantizer which validates the provided quantization config and defines an abstract function which any quantizer needs to implment.
 
         Args:
             quantization_config: quantizer config class contains all the information about the quantizer configuration.
+            freeze_quant_params: whether to freeze all learnable quantization parameters during training.
         """
 
         # verify the quantizer class that inherits this class only has a config argument and key-word arguments
@@ -85,6 +85,7 @@ class BaseTrainableQuantizer(BaseInferableQuantizer, ABC):
                 f"Unrecognized 'QuantizationTarget': {static_quantization_target}.")  # pragma: no cover
 
         self.quantizer_parameters = {}
+        self.freeze_quant_params = freeze_quant_params
 
     @classmethod
     def get_sig(cls):

model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/ste/symmetric_ste.py

@@ -18,7 +18,8 @@ import numpy as np
 import torch
 from torch import nn
 
-from mct_quantizers import mark_quantizer, QuantizationTarget, QuantizationMethod, PytorchQuantizationWrapper
+from mct_quantizers import mark_quantizer, QuantizationTarget, QuantizationMethod, PytorchQuantizationWrapper, \
+    PytorchActivationQuantizationHolder
 from mct_quantizers.pytorch.quantizers import ActivationPOTInferableQuantizer, ActivationSymmetricInferableQuantizer
 from model_compression_toolkit import constants as C
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
@@ -39,14 +40,15 @@ class STESymmetricActivationTrainableQuantizer(BasePytorchActivationTrainableQua
     Trainable constrained quantizer to quantize a layer activations.
     """
 
-    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
+    def __init__(self, quantization_config: TrainableQuantizerActivationConfig, freeze_quant_params: bool = False):
         """
         Initialize a STESymmetricActivationTrainableQuantizer object with parameters to use for symmetric or power of two quantization.
 
         Args:
             quantization_config: trainable quantizer config class
+            freeze_quant_params: whether to freeze learnable quantization parameters
         """
-        super().__init__(quantization_config)
+        super().__init__(quantization_config, freeze_quant_params)
         self.power_of_two = quantization_config.activation_quantization_method == QuantizationMethod.POWER_OF_TWO
         self.sign = quantization_config.activation_quantization_params['is_signed']
         np_threshold_values = quantization_config.activation_quantization_params[C.THRESHOLD]
@@ -56,7 +58,7 @@ class STESymmetricActivationTrainableQuantizer(BasePytorchActivationTrainableQua
     def initialize_quantization(self,
                                 tensor_shape: torch.Size,
                                 name: str,
-                                layer: PytorchQuantizationWrapper):
+                                layer: PytorchActivationQuantizationHolder):
         """
         Add quantizer parameters to the quantizer parameters dictionary
 
@@ -66,7 +68,7 @@ class STESymmetricActivationTrainableQuantizer(BasePytorchActivationTrainableQua
             layer: Layer to quantize.
         """
         layer.register_parameter(name, nn.Parameter(to_torch_tensor(self.threshold_tensor),
-                                                    requires_grad=True))
+                                                    requires_grad=not self.freeze_quant_params))
 
         # save the quantizer added parameters for later calculations
         self.add_quantizer_variable(THRESHOLD_TENSOR, layer.get_parameter(name), VariableGroup.QPARAMS)

model_compression_toolkit/trainable_infrastructure/pytorch/activation_quantizers/ste/uniform_ste.py

@@ -36,14 +36,15 @@ class STEUniformActivationTrainableQuantizer(BasePytorchActivationTrainableQuant
     Trainable constrained quantizer to quantize a layer activations.
     """
 
-    def __init__(self, quantization_config: TrainableQuantizerActivationConfig):
+    def __init__(self, quantization_config: TrainableQuantizerActivationConfig, freeze_quant_params: bool = False):
         """
         Initialize a STEUniformActivationTrainableQuantizer object with parameters to use for uniform quantization.
 
         Args:
-            quantization_config: trainable quantizer config class
+            quantization_config: trainable quantizer config class.
+            freeze_quant_params: whether to freeze learnable quantization parameters.
         """
-        super().__init__(quantization_config)
+        super().__init__(quantization_config, freeze_quant_params)
 
         np_min_range = quantization_config.activation_quantization_params[C.RANGE_MIN]
         np_max_range = quantization_config.activation_quantization_params[C.RANGE_MAX]
@@ -56,7 +57,7 @@ class STEUniformActivationTrainableQuantizer(BasePytorchActivationTrainableQuant
                                 name: str,
                                 layer: PytorchQuantizationWrapper):
         """
-        Add quantizer parameters to the quantizer parameters dictionary
+        Add quantizer parameters to the quantizer parameters dictionary.
 
         Args:
             tensor_shape: tensor shape of the quantized tensor.
@@ -64,9 +65,9 @@ class STEUniformActivationTrainableQuantizer(BasePytorchActivationTrainableQuant
             layer: Layer to quantize.
         """
         layer.register_parameter(name+"_"+FQ_MIN, nn.Parameter(to_torch_tensor(self.min_range_tensor),
-                                                               requires_grad=True))
+                                                               requires_grad=not self.freeze_quant_params))
         layer.register_parameter(name+"_"+FQ_MAX, nn.Parameter(to_torch_tensor(self.max_range_tensor),
-                                                               requires_grad=True))
+                                                               requires_grad=not self.freeze_quant_params))
 
         # Save the quantizer parameters for later calculations
         self.add_quantizer_variable(FQ_MIN, layer.get_parameter(name+"_"+FQ_MIN), VariableGroup.QPARAMS)

model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py

@@ -46,6 +46,14 @@ if FOUND_TORCH:
                 quantizer_parameter, parameter_group = parameter_dict[VAR], parameter_dict[GROUP]
                 if quantizer_parameter.requires_grad and parameter_group == group:
                     quantizer_trainable.append(quantizer_parameter)
+
+            # sanity check to catch inconsistent initialization
+            if self.freeze_quant_params and group == VariableGroup.QPARAMS and quantizer_trainable:
+                Logger.critical(
+                    'Found trainable quantization params despite self.freeze_quant_params=True. '
+                    'Quantization parameters were probably not initialized correctly in the Quantizer.'
+                )    # pragma: no cover
+
             return quantizer_trainable
 
 else: