mct-nightly 1.8.0.22032023.post333__py3-none-any.whl → 1.8.0.22052023.post408__py3-none-any.whl

model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py

@@ -12,63 +12,86 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from enum import Enum
 from typing import Callable
 
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import FOUND_TORCH
-
-
-class PyTorchExportMode(Enum):
-    FAKELY_QUANT_TORCHSCRIPT = 0
-    FAKELY_QUANT_ONNX = 1
-
+from model_compression_toolkit.constants import FOUND_TORCH
+from model_compression_toolkit.exporter.model_exporter.pytorch.export_serialization_format import \
+    PytorchExportSerializationFormat
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.target_platform.quantization_format import \
+    QuantizationFormat
 
 if FOUND_TORCH:
     import torch.nn
-    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_onnx_pytorch_exporter import FakelyQuantONNXPyTorchExporter
-    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_torchscript_pytorch_exporter import FakelyQuantTorchScriptPyTorchExporter
+    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_onnx_pytorch_exporter import \
+        FakelyQuantONNXPyTorchExporter
+    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_torchscript_pytorch_exporter import \
+        FakelyQuantTorchScriptPyTorchExporter
     from model_compression_toolkit.exporter.model_wrapper.pytorch.validate_layer import is_pytorch_layer_exportable
 
+    supported_serialization_quantization_export_dict = {
+        PytorchExportSerializationFormat.TORCHSCRIPT: [QuantizationFormat.FAKELY_QUANT],
+        PytorchExportSerializationFormat.ONNX: [QuantizationFormat.FAKELY_QUANT]
+    }
+
     def pytorch_export_model(model: torch.nn.Module,
                              save_model_path: str,
                              repr_dataset: Callable,
+                             target_platform_capabilities: TargetPlatformCapabilities,
                              is_layer_exportable_fn: Callable = is_pytorch_layer_exportable,
-                             mode: PyTorchExportMode = PyTorchExportMode.FAKELY_QUANT_TORCHSCRIPT) -> None:
+                             serialization_format: PytorchExportSerializationFormat =
+                             PytorchExportSerializationFormat.TORCHSCRIPT) -> None:
         """
         Export a PyTorch quantized model to a torchscript or onnx model.
         The model will be saved to the path in save_model_path.
-        Mode can be used for different exported files. Currently, pytorch_export_model
-        supports PyTorchExportMode.FAKELY_QUANT_TORCHSCRIPT (where the exported model
-        is in a TorchScript format and its weights and activations are float fakely-quantized values),
-        and PyTorchExportMode.FakelyQuantONNX (where the exported model
-        is in an ONNX format and its weights and activations are float fakely-quantized values)
+        Currently, pytorch_export_model supports only QuantizationFormat.FAKELY_QUANT (where weights
+        and activations are float fakely-quantized values) and PytorchExportSerializationFormat.TORCHSCRIPT
+        (where the model will be saved to TorchScript model) or PytorchExportSerializationFormat.ONNX
+        (where the model will be saved to ONNX model).
 
         Args:
             model: Model to export.
-            is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
-            mode: Mode to export the model according to.
             save_model_path: Path to save the model.
             repr_dataset: Representative dataset for tracing the pytorch model (mandatory for exporting it).
+            target_platform_capabilities: TargetPlatformCapabilities object that describes the desired inference
+            target platform (includes quantization format).
+            is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
+            serialization_format: Format to export the model according to (by default
+            PytorchExportSerializationFormat.TORCHSCRIPT).
 
         """
 
-        if mode == PyTorchExportMode.FAKELY_QUANT_TORCHSCRIPT:
-            exporter = FakelyQuantTorchScriptPyTorchExporter(model,
-                                                             is_layer_exportable_fn,
-                                                             save_model_path,
-                                                             repr_dataset)
+        if serialization_format == PytorchExportSerializationFormat.TORCHSCRIPT:
+            if target_platform_capabilities.tp_model.quantization_format in \
+                    supported_serialization_quantization_export_dict[serialization_format]:
+                exporter = FakelyQuantTorchScriptPyTorchExporter(model,
+                                                                 is_layer_exportable_fn,
+                                                                 save_model_path,
+                                                                 repr_dataset)
+            else:
+                Logger.critical(
+                    f'Unsupported quantization {target_platform_capabilities.tp_model.quantization_format} for '
+                    f'serialization {serialization_format} was used to export Pytorch model. Please see API for '
+                    f'supported formats.')  # pragma: no cover
 
-        elif mode == PyTorchExportMode.FAKELY_QUANT_ONNX:
-            exporter = FakelyQuantONNXPyTorchExporter(model,
-                                                      is_layer_exportable_fn,
-                                                      save_model_path,
-                                                      repr_dataset)
+        elif serialization_format == PytorchExportSerializationFormat.ONNX:
+            if target_platform_capabilities.tp_model.quantization_format in \
+                    supported_serialization_quantization_export_dict[serialization_format]:
+                exporter = FakelyQuantONNXPyTorchExporter(model,
+                                                          is_layer_exportable_fn,
+                                                          save_model_path,
+                                                          repr_dataset)
+            else:
+                Logger.critical(
+                    f'Unsupported quantization {target_platform_capabilities.tp_model.quantization_format} for '
+                    f'serialization {serialization_format} was used to export Pytorch model. Please see API for '
+                    f'supported formats.')  # pragma: no cover
 
         else:
             Logger.critical(
-                f'Unsupported mode was used {mode.name} to export PyTorch model. '
-                f'Please see API for supported modes.')  # pragma: no cover
+                f'Unsupported serialization {serialization_format} was used to export Pytorch model. Please see API '
+                f'for supported formats.')  # pragma: no cover
 
         exporter.export()
 

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

@@ -17,9 +17,10 @@ from typing import Tuple
 
 from model_compression_toolkit import quantizers_infrastructure as qi
 from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common import Graph, Logger
-from model_compression_toolkit.core.common.constants import FOUND_TF
+from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.constants import FOUND_TF
 from model_compression_toolkit.core.common.user_info import UserInformation
+from model_compression_toolkit.logger import Logger
 
 if FOUND_TF:
     import tensorflow as tf
@@ -34,6 +35,7 @@ if FOUND_TF:
         Args:
             n: A node of mct graph.
             layer: A keras layer
+            include_activation_quantizers: Whether to use the wrapper for the activation quantizer or not
 
         Returns: Wrapped layer with weights quantizers and activation quantizers
 
@@ -55,7 +57,7 @@ if FOUND_TF:
             Exportable Keras model and user information.
         """
         exportable_model, user_info = KerasModelBuilder(graph=graph,
-                                             wrapper=_get_wrapper).build_model()
+                                                        wrapper=_get_wrapper).build_model()
         exportable_model.trainable = False
         return exportable_model, user_info
 else:

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

@@ -14,9 +14,11 @@
 # ==============================================================================
 from typing import Dict, Any
 
-from model_compression_toolkit.core.common import BaseNode, Logger
-from model_compression_toolkit.core.common.constants import THRESHOLD, RANGE_MIN, RANGE_MAX, SIGNED, CLUSTER_CENTERS, SCALE_PER_CHANNEL
-from model_compression_toolkit.core.common.target_platform import QuantizationMethod
+from model_compression_toolkit.core.common import BaseNode
+from model_compression_toolkit.constants import THRESHOLD, RANGE_MIN, RANGE_MAX, SIGNED, CLUSTER_CENTERS, SCALE_PER_CHANNEL
+
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import QuantizationTarget
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.get_quantizers import get_inferable_quantizer_class
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers.base_keras_inferable_quantizer import BaseKerasInferableQuantizer

model_compression_toolkit/exporter/model_wrapper/keras/validate_layer.py

@@ -15,8 +15,8 @@
 from typing import Any
 
 
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import FOUND_TF
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import FOUND_TF
 
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.base_inferable_quantizer import BaseInferableQuantizer
 

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

@@ -16,8 +16,9 @@
 
 from model_compression_toolkit import quantizers_infrastructure as qi
 from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common import Graph, Logger
-from model_compression_toolkit.core.common.constants import FOUND_TORCH
+from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.constants import FOUND_TORCH
+from model_compression_toolkit.logger import Logger
 
 if FOUND_TORCH:
     import torch

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

@@ -15,10 +15,11 @@
 
 from typing import Dict, Any
 
-from model_compression_toolkit.core.common import BaseNode, Logger
-from model_compression_toolkit.core.common.constants import THRESHOLD, SIGNED, RANGE_MIN, RANGE_MAX, \
+from model_compression_toolkit.core.common import BaseNode
+from model_compression_toolkit.constants import THRESHOLD, SIGNED, RANGE_MIN, RANGE_MAX, \
     SCALE_PER_CHANNEL, CLUSTER_CENTERS
-from model_compression_toolkit.core.common.target_platform import QuantizationMethod
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
 from model_compression_toolkit.quantizers_infrastructure import QuantizationTarget
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.get_quantizers import \
     get_inferable_quantizer_class

model_compression_toolkit/exporter/model_wrapper/pytorch/validate_layer.py

@@ -14,8 +14,8 @@
 # ==============================================================================
 from typing import Any
 
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import FOUND_TORCH
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import FOUND_TORCH
 
 if FOUND_TORCH:
     from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper

model_compression_toolkit/gptq/common/gptq_config.py

@@ -111,10 +111,8 @@ class GradientPTQConfig:
         self.regularization_factor = regularization_factor
         self.hessian_weights_config = hessian_weights_config
 
-        # Since the default quantizer is soft quantizer, we initialize the gptq_quantizer_params_override dictionary
-        # with its extended params
-        self.gptq_quantizer_params_override = {QUANT_PARAM_LEARNING_STR: False} \
-            if gptq_quantizer_params_override is None else gptq_quantizer_params_override
+        self.gptq_quantizer_params_override = {} if gptq_quantizer_params_override is None \
+            else gptq_quantizer_params_override
 
 
 class GradientPTQConfigV2(GradientPTQConfig):

model_compression_toolkit/gptq/common/gptq_framework_implementation.py

@@ -0,0 +1,32 @@
+# Copyright 2023 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 abc import abstractmethod
+
+from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
+
+
+class GPTQFrameworkImplemantation(FrameworkImplementation):
+    """
+    Class to implement framework related methods that are used in GPTQ
+    """
+
+    @abstractmethod
+    def get_gptq_trainer_obj(self):
+        """
+        Returns: GPTQTrainer object
+        """
+        raise NotImplemented(f'{self.__class__.__name__} have to implement the '
+                             f'framework\'s get_gptq_trainer method.')  # pragma: no cover

model_compression_toolkit/gptq/common/gptq_graph.py

@@ -14,8 +14,8 @@
 # ==============================================================================
 from typing import Tuple, List
 
-from model_compression_toolkit import FrameworkInfo
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.core import FrameworkInfo
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 

model_compression_toolkit/gptq/common/gptq_training.py

@@ -17,12 +17,13 @@ from abc import ABC, abstractmethod
 import numpy as np
 from typing import Callable, List, Any
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig
-from model_compression_toolkit.core.common import Graph, Logger, BaseNode
+from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.gptq.common.gptq_constants import QUANT_PARAM_LEARNING_STR
+from model_compression_toolkit.gptq.common.gptq_framework_implementation import GPTQFrameworkImplemantation
 from model_compression_toolkit.gptq.common.gptq_graph import get_compare_points
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
+from model_compression_toolkit.logger import Logger
 
 
 class GPTQTrainer(ABC):
@@ -34,7 +35,7 @@ class GPTQTrainer(ABC):
                  graph_float: Graph,
                  graph_quant: Graph,
                  gptq_config: GradientPTQConfig,
-                 fw_impl: FrameworkImplementation,
+                 fw_impl: GPTQFrameworkImplemantation,
                  fw_info: FrameworkInfo):
         """
         Build two models from a graph: A teacher network (float model) and a student network (quantized model).
@@ -259,7 +260,7 @@ def gptq_training(graph_float: Graph,
                   graph_quant: Graph,
                   gptq_config: GradientPTQConfig,
                   representative_data_gen: Callable,
-                  fw_impl: FrameworkImplementation,
+                  fw_impl: GPTQFrameworkImplemantation,
                   fw_info: FrameworkInfo) -> Graph:
     """
     GPTQ training process using knowledge distillation with a teacher network (float model) and a student network (quantized model).

model_compression_toolkit/gptq/keras/gptq_keras_implementation.py

@@ -0,0 +1,29 @@
+# Copyright 2023 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 Type
+
+from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
+from model_compression_toolkit.gptq.common.gptq_framework_implementation import GPTQFrameworkImplemantation
+from model_compression_toolkit.gptq.keras.gptq_training import KerasGPTQTrainer
+
+
+class GPTQKerasImplemantation(GPTQFrameworkImplemantation, KerasImplementation):
+
+    def get_gptq_trainer_obj(self) -> Type[KerasGPTQTrainer]:
+        """
+        Returns:  Keras object of GPTQTrainer
+        """
+        return KerasGPTQTrainer

model_compression_toolkit/gptq/keras/gptq_training.py

@@ -16,17 +16,18 @@ from typing import Callable, List, Tuple, Union
 
 import tensorflow as tf
 from keras import Model
+from packaging import version
 from tensorflow.keras.layers import Layer
 from tqdm import tqdm
 
 # As from Tensorflow 2.6, keras is a separate package and some classes should be imported differently.
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
-from packaging import version
-
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
 from model_compression_toolkit.gptq.keras.quantizer.quantization_builder import quantization_builder
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.activation_quantization_holder import ActivationQuantizationHolder
 
 if version.parse(tf.__version__) < version.parse("2.6"):
     from tensorflow.python.keras.engine.base_layer import TensorFlowOpLayer
@@ -105,7 +106,7 @@ class KerasGPTQTrainer(GPTQTrainer):
             [len(optimizer_params_tuple[1]) for optimizer_params_tuple in self.optimizer_with_param]) > 0
 
         if self.float_user_info.input_scale != self.gptq_user_info.input_scale:
-            common.Logger.error("Input scale mismatch between float and GPTQ networks")  # pragma: no cover
+            Logger.error("Input scale mismatch between float and GPTQ networks")  # pragma: no cover
         else:
             self.input_scale = self.gptq_user_info.input_scale
 
@@ -113,8 +114,8 @@ class KerasGPTQTrainer(GPTQTrainer):
 
         self.reg_func = get_regularization(self.gptq_config, representative_data_gen)
 
-    def _is_gptq_applicable(self,
-                            node: common.BaseNode) -> bool:
+    def _is_gptq_weights_trainable(self,
+                                   node: common.BaseNode) -> bool:
         """
         A function for deciding if a layer should be fine-tuned during GPTQ.
 
@@ -126,11 +127,13 @@ class KerasGPTQTrainer(GPTQTrainer):
         """
 
         if node.is_weights_quantization_enabled() and not self.fw_info.is_kernel_op(node.type):
-            common.Logger.error(f"GPTQ Error: Quantizing node {node.name} of type {node.type} "
+            Logger.error(f"GPTQ Error: Quantizing node {node.name} of type {node.type} "
                                 f"without a kernel isn't supported")
         return node.is_weights_quantization_enabled()
 
-    def gptq_wrapper(self, n: common.BaseNode, layer: Layer) -> Union[qi.KerasQuantizationWrapper, Layer]:
+    def gptq_wrapper(self,
+                     n: common.BaseNode,
+                     layer: Layer) -> Union[qi.KerasQuantizationWrapper, Layer]:
         """
         A function which takes a computational graph node and a keras layer and perform the quantization wrapping.
 
@@ -141,14 +144,37 @@ class KerasGPTQTrainer(GPTQTrainer):
         Returns: Wrapped layer if the layer should be wrap, otherwise returns the layer as is.
 
         """
-        if self._is_gptq_applicable(n):
-            weights_quantizers, activation_quantizers = quantization_builder(n, self.gptq_config)
+        if self._is_gptq_weights_trainable(n):
+            weights_quantizers, _ = quantization_builder(n, self.gptq_config) # TODO: split quantizers building into two functions: for weights and activations
             return qi.KerasQuantizationWrapper(layer,
-                                               weights_quantizers=weights_quantizers,
-                                               activation_quantizers=activation_quantizers)
+                                               weights_quantizers=weights_quantizers)
         else:
             return layer
 
+    def get_activation_quantizer_holder(self, n: common.BaseNode) -> Union[None, Callable]:
+        """
+        Retrieve a ActivationQuantizationHolder layer to use for activation quantization for a node.
+        If the layer is not supposed to be wrapped with activation quantizers - return None.
+
+        Args:
+            n: Node to get ActivationQuantizationHolder to attach in its output.
+
+        Returns:
+            A ActivationQuantizationHolder layer for the node activation quantization.
+        """
+        _, activation_quantizers = quantization_builder(n, self.gptq_config) # TODO: split quantizers building into two functions: for weights and activations
+
+        # Holder by definition uses a single quantizer for the activation quantization
+        # thus we make sure this is the only possible case (unless it's a node with no activation
+        # quantization, which in this case has an empty list).
+        if len(activation_quantizers) == 1:
+            return ActivationQuantizationHolder(activation_quantizers[0])
+
+        Logger.error(
+            f'ActivationQuantizationHolder supports a single quantizer but {len(activation_quantizers)} quantizers '
+            f'were found for node {n}')
+
+
     def build_gptq_model(self) -> Tuple[Model, UserInformation]:
         """
         Build the GPTQ model with QuantizationWrappers
@@ -161,7 +187,8 @@ class KerasGPTQTrainer(GPTQTrainer):
                                                        append2output=self.compare_points,
                                                        fw_info=self.fw_info,
                                                        return_float_outputs=True,
-                                                       wrapper=self.gptq_wrapper).build_model()
+                                                       wrapper=self.gptq_wrapper,
+                                                       get_activation_quantizer_holder_fn=self.get_activation_quantizer_holder).build_model()
 
         return gptq_model, gptq_user_info
 
@@ -280,7 +307,7 @@ class KerasGPTQTrainer(GPTQTrainer):
                     self.gptq_config.log_function(loss_value_step, grads[0], in_optimizer_with_param[0][-1],
                                                   self.compare_points)
                 self.loss_list.append(loss_value_step.numpy())
-                common.Logger.debug(f'last loss value: {self.loss_list[-1]}')
+                Logger.debug(f'last loss value: {self.loss_list[-1]}')
 
     def update_graph(self):
         """
@@ -297,7 +324,7 @@ class KerasGPTQTrainer(GPTQTrainer):
                 if len(node) == 0 and isinstance(layer.layer, TensorFlowOpLayer):
                     node = graph.find_node_by_name('_'.join(layer.layer.name.split('_')[3:]))
                 if len(node) != 1:
-                    common.Logger.error(f"Can't update GPTQ graph due to missing layer named: {layer.layer.name}")
+                    Logger.error(f"Can't update GPTQ graph due to missing layer named: {layer.layer.name}")
                 node = node[0]
                 kernel_attribute = get_kernel_attribute_name_for_gptq(layer_type=node.type,
                                                                       fw_info=self.fw_info)

model_compression_toolkit/gptq/keras/graph_info.py

@@ -20,6 +20,7 @@ from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from tensorflow.keras.models import Model
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.gptq.common.gptq_graph import get_kernel_attribute_name_for_gptq
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper
 from model_compression_toolkit.quantizers_infrastructure.trainable_infrastructure.common.base_trainable_quantizer import VariableGroup
 
@@ -50,6 +51,9 @@ def get_gptq_trainable_parameters(fxp_model: Model,
                                                                   fw_info=DEFAULT_KERAS_INFO)
 
             # collect trainable weights per quantizer
+            if kernel_attribute not in layer.weights_quantizers:
+                Logger.error(f'{kernel_attribute} was not found in weight quantizers of layer {layer.layer}')
+
             quantizer_trainable_weights = layer.weights_quantizers[kernel_attribute].get_trainable_variables(VariableGroup.WEIGHTS)
             quantizer_trainable_threshold = layer.weights_quantizers[kernel_attribute].get_trainable_variables(VariableGroup.QPARAMS)
             trainable_weights.append(quantizer_trainable_weights)

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -16,21 +16,19 @@
 from typing import Callable, Tuple
 from packaging import version
 
-from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import TENSORFLOW
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import TENSORFLOW, FOUND_TF
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
-    MixedPrecisionQuantizationConfigV2
-from model_compression_toolkit import CoreConfig
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.runner import core_runner, _init_tensorboard_writer
 from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.core.exporter import export_model
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
-from model_compression_toolkit.core.common.target_platform.targetplatform2framework import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
 
 LR_DEFAULT = 0.15
 LR_REST_DEFAULT = 1e-4
@@ -38,14 +36,14 @@ LR_BIAS_DEFAULT = 1e-4
 LR_QUANTIZATION_PARAM_DEFAULT = 1e-3
 GPTQ_MOMENTUM = 0.9
 
-if common.constants.FOUND_TF:
+if FOUND_TF:
     import tensorflow as tf
     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.gptq.keras.gptq_keras_implementation import GPTQKerasImplemantation
     from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
     from model_compression_toolkit.gptq.keras.gptq_loss import GPTQMultipleTensorsLoss
-    from model_compression_toolkit.core.keras.constants import DEFAULT_TP_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
     from model_compression_toolkit import get_target_platform_capabilities
 
@@ -62,7 +60,8 @@ if common.constants.FOUND_TF:
                               optimizer: OptimizerV2 = tf.keras.optimizers.Adam(learning_rate=LR_DEFAULT),
                               optimizer_rest: OptimizerV2 = tf.keras.optimizers.Adam(learning_rate=LR_REST_DEFAULT),
                               loss: Callable = GPTQMultipleTensorsLoss(),
-                              log_function: Callable = None) -> GradientPTQConfigV2:
+                              log_function: Callable = None,
+                              use_hessian_based_weights: bool = True) -> GradientPTQConfigV2:
         """
         Create a GradientPTQConfigV2 instance for Keras models.
 
@@ -72,6 +71,7 @@ if common.constants.FOUND_TF:
             optimizer_rest (OptimizerV2): Keras optimizer to use for fine-tuning of the bias variable.
             loss (Callable): loss to use during fine-tuning. should accept 4 lists of tensors. 1st list of quantized tensors, the 2nd list is the float tensors, the 3rd is a list of quantized weights and the 4th is a list of float weights.
             log_function (Callable): Function to log information about the gptq process.
+            use_hessian_based_weights (bool): Whether to use Hessian-based weights for weighted average loss.
 
         returns:
             a GradientPTQConfigV2 object to use when fine-tuning the quantized model using gptq.
@@ -94,9 +94,16 @@ if common.constants.FOUND_TF:
             The configuration can be passed to :func:`~model_compression_toolkit.keras_post_training_quantization` in order to quantize a keras model using gptq.
 
         """
-        bias_optimizer = tf.keras.optimizers.SGD(learning_rate=LR_BIAS_DEFAULT, momentum=GPTQ_MOMENTUM)
-        return GradientPTQConfigV2(n_epochs, optimizer, optimizer_rest=optimizer_rest, loss=loss,
-                                   log_function=log_function, train_bias=True, optimizer_bias=bias_optimizer)
+        bias_optimizer = tf.keras.optimizers.SGD(learning_rate=LR_BIAS_DEFAULT,
+                                                 momentum=GPTQ_MOMENTUM)
+        return GradientPTQConfigV2(n_epochs,
+                                   optimizer,
+                                   optimizer_rest=optimizer_rest,
+                                   loss=loss,
+                                   log_function=log_function,
+                                   train_bias=True,
+                                   optimizer_bias=bias_optimizer,
+                                   use_hessian_based_weights=use_hessian_based_weights)
 
 
     def keras_gradient_post_training_quantization_experimental(in_model: Model,
@@ -158,20 +165,20 @@ if common.constants.FOUND_TF:
 
             Create an MCT core config, containing the quantization configuration:
 
-            >>> config = mct.CoreConfig()
+            >>> config = mct.core.CoreConfig()
 
             If mixed precision is desired, create an MCT core config with a mixed-precision configuration, to quantize a model
             with different bitwidths for different layers.
             The candidates bitwidth for quantization should be defined in the target platform model:
 
-            >>> config = mct.CoreConfig(mixed_precision_config=mct.MixedPrecisionQuantizationConfigV2(num_of_images=1))
+            >>> config = mct.core.CoreConfig(mixed_precision_config=mct.core.MixedPrecisionQuantizationConfigV2(num_of_images=1))
 
             For mixed-precision set a target KPI object:
             Create a KPI object to limit our returned model's size. Note that this value affects only coefficients
             that should be quantized (for example, the kernel of Conv2D in Keras will be affected by this value,
             while the bias will not):
 
-            >>> kpi = mct.KPI(model.count_params() * 0.75)  # About 0.75 of the model size when quantized with 8 bits.
+            >>> kpi = mct.core.KPI(model.count_params() * 0.75)  # About 0.75 of the model size when quantized with 8 bits.
 
             Create GPTQ config:
 
@@ -187,15 +194,15 @@ if common.constants.FOUND_TF:
 
         if core_config.mixed_precision_enable:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
-                common.Logger.error("Given quantization config to mixed-precision facade is not of type "
+                Logger.error("Given quantization config to mixed-precision facade is not of type "
                                     "MixedPrecisionQuantizationConfigV2. Please use keras_post_training_quantization "
                                     "API, or pass a valid mixed precision configuration.")  # pragma: no cover
 
-            common.Logger.info("Using experimental mixed-precision quantization. "
+            Logger.info("Using experimental mixed-precision quantization. "
                                "If you encounter an issue please file a bug.")
         tb_w = _init_tensorboard_writer(fw_info)
 
-        fw_impl = KerasImplementation()
+        fw_impl = GPTQKerasImplemantation()
 
         tg, bit_widths_config = core_runner(in_model=in_model,
                                             representative_data_gen=representative_data_gen,

model_compression_toolkit/gptq/keras/quantizer/__init__.py

@@ -15,3 +15,4 @@
 
 import model_compression_toolkit.gptq.keras.quantizer.ste_rounding.symmetric_ste
 import model_compression_toolkit.gptq.keras.quantizer.soft_rounding.symmetric_soft_quantizer
+import model_compression_toolkit.gptq.keras.quantizer.soft_rounding.uniform_soft_quantizer

model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py

@@ -15,8 +15,8 @@
 from abc import abstractmethod
 from typing import Union, Dict, List
 
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import FOUND_TF
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import FOUND_TF
 from model_compression_toolkit.gptq.common.gptq_constants import WEIGHTS_QUANTIZATION_PARAMS
 
 from model_compression_toolkit.quantizers_infrastructure import TrainableQuantizerWeightsConfig, \