mct-nightly 2.2.0.20240908.528__tar.gz → 2.2.0.20240909.453__tar.gz

{mct-nightly-2.2.0.20240908.528 → mct-nightly-2.2.0.20240909.453}/PKG-INFO

@@ -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 → mct-nightly-2.2.0.20240909.453}/mct_nightly.egg-info/PKG-INFO

@@ -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 → mct-nightly-2.2.0.20240909.453}/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"

{mct-nightly-2.2.0.20240908.528 → mct-nightly-2.2.0.20240909.453}/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

{mct-nightly-2.2.0.20240908.528 → mct-nightly-2.2.0.20240909.453}/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):

{mct-nightly-2.2.0.20240908.528 → mct-nightly-2.2.0.20240909.453}/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)

{mct-nightly-2.2.0.20240908.528 → mct-nightly-2.2.0.20240909.453}/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)

{mct-nightly-2.2.0.20240908.528 → mct-nightly-2.2.0.20240909.453}/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: