mct-nightly 1.8.0.22042023.post414__py3-none-any.whl → 1.8.0.22052023.post408__py3-none-any.whl

model_compression_toolkit/exporter/model_exporter/{tflite → keras}/fakely_quant_tflite_exporter.py

@@ -20,7 +20,7 @@ import keras.models
 import tensorflow as tf
 
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.load_model import keras_load_quantized_model
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.exporter.model_exporter.keras.fakely_quant_keras_exporter import FakelyQuantKerasExporter
 
 

model_compression_toolkit/exporter/model_exporter/{tflite → keras}/int8_tflite_exporter.py

@@ -23,7 +23,7 @@ from keras.layers import Dense, Conv2D, Reshape
 from keras.models import clone_model
 
 from model_compression_toolkit import quantizers_infrastructure as qi
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.exporter.model_exporter.keras.fakely_quant_keras_exporter import FakelyQuantKerasExporter
 from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers import \
     constants as keras_inferable_constants

model_compression_toolkit/exporter/model_exporter/keras/keras_export_facade.py

@@ -12,53 +12,91 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from enum import Enum
 from typing import Callable, Dict
 
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import FOUND_TF
-
-
-class KerasExportMode(Enum):
-    FAKELY_QUANT = 0
-
+from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.exporter.model_exporter.keras.export_serialization_format import \
+    KerasExportSerializationFormat
+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_TF:
     import keras
     from model_compression_toolkit.exporter.model_wrapper.keras.validate_layer import is_keras_layer_exportable
-    from model_compression_toolkit.exporter.model_exporter.keras.fakely_quant_keras_exporter import FakelyQuantKerasExporter
+    from model_compression_toolkit.exporter.model_exporter.keras.fakely_quant_keras_exporter import \
+        FakelyQuantKerasExporter
+    from model_compression_toolkit.exporter.model_exporter.keras.fakely_quant_tflite_exporter import \
+        FakelyQuantTFLiteExporter
+    from model_compression_toolkit.exporter.model_exporter.keras.int8_tflite_exporter import INT8TFLiteExporter
+
+    supported_serialization_quantization_export_dict = {
+        KerasExportSerializationFormat.KERAS_H5: [QuantizationFormat.FAKELY_QUANT],
+        KerasExportSerializationFormat.TFLITE: [QuantizationFormat.FAKELY_QUANT, QuantizationFormat.INT8]
+    }
 
     def keras_export_model(model: keras.models.Model,
                            save_model_path: str,
+                           target_platform_capabilities: TargetPlatformCapabilities,
                            is_layer_exportable_fn: Callable = is_keras_layer_exportable,
-                           mode: KerasExportMode = KerasExportMode.FAKELY_QUANT) -> Dict[str, type]:
+                           serialization_format: KerasExportSerializationFormat =
+                           KerasExportSerializationFormat.KERAS_H5) -> \
+            Dict[str, type]:
         """
-        Export a Keras quantized model to h5 model.
+        Export a Keras quantized model to a h5 or tflite model.
         The model will be saved to the path in save_model_path.
-        Mode can be used for different exported files. Currently, keras_export_model
-        supports KerasExportMode.FAKELY_QUANT (where weights and activations are
-        float fakely-quantized values).
+        keras_export_model supports the combination of QuantizationFormat.FAKELY_QUANT (where weights
+        and activations are float fakely-quantized values) and KerasExportSerializationFormat.KERAS_H5 (where the model
+        will be saved to h5 model) or the combination of KerasExportSerializationFormat.TFLITE (where the model will be
+        saved to tflite model) with QuantizationFormat.FAKELY_QUANT or QuantizationFormat.INT8 (where weights and
+        activations are represented using 8bits integers).
 
         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.
+            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
+            KerasExportSerializationFormat.KERAS_H5).
 
         Returns:
             Custom objects dictionary needed to load the model.
 
         """
 
-        if mode == KerasExportMode.FAKELY_QUANT:
-            exporter = FakelyQuantKerasExporter(model,
-                                                is_layer_exportable_fn,
-                                                save_model_path)
+        if serialization_format == KerasExportSerializationFormat.KERAS_H5:
+            if target_platform_capabilities.tp_model.quantization_format == QuantizationFormat.FAKELY_QUANT:
+                exporter = FakelyQuantKerasExporter(model,
+                                                    is_layer_exportable_fn,
+                                                    save_model_path)
+            else:
+                Logger.critical(
+                    f'Unsupported quantization {target_platform_capabilities.tp_model.quantization_format} for '
+                    f'serialization {serialization_format} was used to export Keras model. Please see API for '
+                    f'supported formats.')  # pragma: no cover
+
+        elif serialization_format == KerasExportSerializationFormat.TFLITE:
+            if target_platform_capabilities.tp_model.quantization_format == QuantizationFormat.FAKELY_QUANT:
+                exporter = FakelyQuantTFLiteExporter(model,
+                                                     is_layer_exportable_fn,
+                                                     save_model_path)
+
+            elif target_platform_capabilities.tp_model.quantization_format == QuantizationFormat.INT8:
+                exporter = INT8TFLiteExporter(model,
+                                              is_layer_exportable_fn,
+                                              save_model_path)
+            else:
+                Logger.critical(
+                    f'Unsupported quantization {target_platform_capabilities.tp_model.quantization_format} for '
+                    f'serialization {serialization_format} was used to export Keras model. Please see API for '
+                    f'supported formats.')  # pragma: no cover
 
         else:
             Logger.critical(
-                f'Unsupported mode was used {mode.name} to '
-                f'export Keras model. Please see API for supported modes.')  # pragma: no cover
+                f'Unsupported serialization {serialization_format} was used to export Keras model. Please see API '
+                f'for supported formats.')  # pragma: no cover
 
         exporter.export()
 

model_compression_toolkit/exporter/model_exporter/pytorch/export_serialization_format.py

@@ -0,0 +1,20 @@
+# 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 enum import Enum
+
+
+class PytorchExportSerializationFormat(Enum):
+    TORCHSCRIPT = 0
+    ONNX = 1

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -16,17 +16,21 @@ from typing import Callable
 
 import torch.nn
 
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from model_compression_toolkit.exporter.model_exporter.pytorch.base_pytorch_exporter import BasePyTorchExporter
 from packaging import version
 
+from model_compression_toolkit.quantizers_infrastructure import PytorchQuantizationWrapper
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.common.constants import LAYER
+
 # ONNX opset version 16 is supported from PyTorch 1.12
 if version.parse(torch.__version__) < version.parse("1.12"):
     OPSET_VERSION = 15
 else:
     OPSET_VERSION = 16
 
+
 class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
     """
     Exporter for fakely-quant PyTorch models.
@@ -70,6 +74,16 @@ class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
 
         Logger.info(f"Exporting PyTorch fake quant onnx model: {self.save_model_path}")
 
+        # Replace float weight with wrapped quantized weights
+        for layer in self.model.modules():
+            if isinstance(layer, PytorchQuantizationWrapper):
+                for name in layer.weights_quantizers.keys():
+                    quantized_weight = torch.nn.Parameter(layer.get_quantized_weights()[name]).detach()
+                    linear_layer = getattr(layer, LAYER)
+                    delattr(linear_layer, name)
+                    setattr(linear_layer, name, torch.nn.Parameter(quantized_weight))
+                layer.weights_quantizers = {}
+
         torch.onnx.export(self.model,
                           model_input,
                           self.save_model_path,

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_torchscript_pytorch_exporter.py

@@ -16,7 +16,7 @@ from typing import Callable
 
 import torch.nn
 
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from model_compression_toolkit.exporter.model_exporter.pytorch.base_pytorch_exporter import BasePyTorchExporter
 

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,8 +14,10 @@
 # ==============================================================================
 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 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

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,9 +15,10 @@
 
 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.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 \

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_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