mct-nightly 1.8.0.8032023.post421__py3-none-any.whl → 1.8.0.8052023.post414__py3-none-any.whl

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/residual_collapsing.py

@@ -20,7 +20,7 @@ from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.core.common.substitutions.residual_collapsing import ResidualCollapsing
 from model_compression_toolkit.core.pytorch.constants import KERNEL
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 
 
 def residual_collapsing_node_matchers() -> Tuple[NodeOperationMatcher, NodeOperationMatcher]:

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/shift_negative_activation.py

@@ -21,7 +21,7 @@ from torch import reshape
 from torch.nn.functional import hardswish, silu, prelu, elu
 from torch.nn.functional import avg_pool2d
 
-from model_compression_toolkit import CoreConfig, FrameworkInfo
+from model_compression_toolkit.core import CoreConfig, FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import EdgeMatcher

model_compression_toolkit/core/pytorch/kpi_data_facade.py

@@ -15,21 +15,21 @@
 
 from typing import Callable
 
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import PYTORCH
-from model_compression_toolkit.core.common.target_platform import TargetPlatformCapabilities
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import PYTORCH
+from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
 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.kpi_tools.kpi_data import compute_kpi_data
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfig, DEFAULT_MIXEDPRECISION_CONFIG, MixedPrecisionQuantizationConfigV2
-from model_compression_toolkit.core.common.constants import FOUND_TORCH
+from model_compression_toolkit.constants import FOUND_TORCH
 
 if FOUND_TORCH:
     from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
     from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
-    from model_compression_toolkit.core.pytorch.constants import DEFAULT_TP_MODEL
+    from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from torch.nn import Module
 
     from model_compression_toolkit import get_target_platform_capabilities
@@ -51,7 +51,7 @@ if FOUND_TORCH:
             representative_data_gen (Callable): Dataset used for calibration.
             quant_config (MixedPrecisionQuantizationConfig): MixedPrecisionQuantizationConfig containing parameters of how the model should be quantized.
             fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices, groups of layers by how they should be quantized, etc.). `Default PyTorch info <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/pytorch/default_framework_info.py>`_
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to. `Default PyTorch TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/tpc_models/pytorch_tp_models/pytorch_default.py>`_
+            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
 
         Returns:
             A KPI object with total weights parameters sum, max activation tensor and total kpi.
@@ -75,7 +75,7 @@ if FOUND_TORCH:
             Import mct and call for KPI data calculation:
 
             >>> import model_compression_toolkit as mct
-            >>> kpi_data = mct.pytorch_kpi_data(module, repr_datagen)
+            >>> kpi_data = mct.core.pytorch_kpi_data(module, repr_datagen)
 
         """
 
@@ -111,7 +111,7 @@ if FOUND_TORCH:
             representative_data_gen (Callable): Dataset used for calibration.
             core_config (CoreConfig): CoreConfig containing parameters for quantization and mixed precision
             fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices, groups of layers by how they should be quantized, etc.). `Default PyTorch info <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/pytorch/default_framework_info.py>`_
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the PyTorch model according to. `Default PyTorch TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/tpc_models/pytorch_tp_models/pytorch_default.py>`_
+            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the PyTorch model according to.
 
         Returns:
 
@@ -132,7 +132,7 @@ if FOUND_TORCH:
             Import mct and call for KPI data calculation:
 
             >>> import model_compression_toolkit as mct
-            >>> kpi_data = mct.pytorch_kpi_data(module, repr_datagen)
+            >>> kpi_data = mct.core.pytorch_kpi_data(module, repr_datagen)
 
         """
 

model_compression_toolkit/core/pytorch/mixed_precision/mixed_precision_wrapper.py

@@ -18,7 +18,7 @@ from typing import Any, List
 import torch
 import copy
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.pytorch.utils import set_model, to_torch_tensor

model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import operator
+from copy import deepcopy
 from typing import List, Any, Tuple, Callable, Type, Dict
 
 import numpy as np
@@ -22,7 +23,7 @@ from torch.nn import Conv2d, ConvTranspose2d, Linear
 from torch.nn import Module, Sigmoid, Softmax
 
 import model_compression_toolkit.core.pytorch.constants as pytorch_constants
-from model_compression_toolkit import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfigV2
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.collectors.statistics_collector import BaseStatsCollector
@@ -74,10 +75,7 @@ from model_compression_toolkit.core.pytorch.pytorch_node_prior_info import creat
 from model_compression_toolkit.core.pytorch.reader.reader import model_reader
 from model_compression_toolkit.core.pytorch.statistics_correction.apply_second_moment_correction import \
     pytorch_apply_second_moment_correction
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
-from model_compression_toolkit.core.pytorch.utils import torch_tensor_to_numpy
-from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
-from model_compression_toolkit.gptq.pytorch.gptq_training import PytorchGPTQTrainer
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor, torch_tensor_to_numpy, set_model
 
 
 class PytorchImplementation(FrameworkImplementation):
@@ -127,7 +125,9 @@ class PytorchImplementation(FrameworkImplementation):
         Returns:
             Graph representing the input module.
         """
-        return model_reader(module, representative_data_gen, self.to_numpy, self.to_tensor)
+        _module = deepcopy(module)
+        _module.eval()
+        return model_reader(_module, representative_data_gen, self.to_numpy, self.to_tensor)
 
     def model_builder(self,
                       graph: Graph,
@@ -323,12 +323,6 @@ class PytorchImplementation(FrameworkImplementation):
             substitutions_list.append(pytorch_batchnorm_refusing())
         return substitutions_list
 
-    def get_gptq_trainer_obj(self) -> Type[GPTQTrainer]:
-        """
-        Returns: GPTQTrainer object
-        """
-        return PytorchGPTQTrainer
-
     def get_sensitivity_evaluator(self,
                                   graph: Graph,
                                   quant_config: MixedPrecisionQuantizationConfigV2,

model_compression_toolkit/core/pytorch/pytorch_node_prior_info.py

@@ -16,7 +16,7 @@ from typing import Any, Tuple
 import numpy as np
 from torch.nn import BatchNorm2d
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.node_prior_info import NodePriorInfo
 from model_compression_toolkit.core.pytorch.constants import MOVING_MEAN, MOVING_VARIANCE, GAMMA, BETA

model_compression_toolkit/core/pytorch/quantizer/fake_quant_builder.py

@@ -12,10 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Tuple, Callable
+from typing import Callable
 import torch
 
-from model_compression_toolkit.core.common.constants import THRESHOLD, SIGNED, RANGE_MIN, RANGE_MAX
+from model_compression_toolkit.constants import THRESHOLD, SIGNED, RANGE_MIN, RANGE_MAX
 from model_compression_toolkit.core.common.quantization.quantizers.uniform_quantizers import threshold_is_power_of_two
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import fix_range_to_include_zero
 

model_compression_toolkit/core/pytorch/quantizer/lut_fake_quant.py

@@ -3,7 +3,7 @@ from typing import Dict, Callable
 import torch
 import numpy as np
 
-from model_compression_toolkit.core.common.constants import SIGNED, CLUSTER_CENTERS, THRESHOLD, MULTIPLIER_N_BITS, EPS
+from model_compression_toolkit.constants import SIGNED, CLUSTER_CENTERS, THRESHOLD, MULTIPLIER_N_BITS, EPS
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 
 

model_compression_toolkit/core/pytorch/reader/graph_builders.py

@@ -25,6 +25,7 @@ from model_compression_toolkit.core.common.graph.functional_node import Function
 from model_compression_toolkit.core.pytorch.constants import OUTPUT, PLACEHOLDER, TENSOR_META, CALL_FUNCTION, TYPE, \
     CALL_METHOD, BIAS, FUNCTIONAL_OP, OP_CALL_KWARGS, OP_CALL_ARGS, INPUTS_AS_LIST, GET_ATTR, CONSTANT, BUFFER
 from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlaceHolder, ConstantHolder, BufferHolder
+from model_compression_toolkit.logger import Logger
 
 
 def extract_holder_weights(constant_name, node_target, model, weights, to_numpy):
@@ -64,6 +65,7 @@ def nodes_builder(model: GraphModule,
     Args:
         model: Pytorch FX model.
         module_dict: A dictionary of the Pyotrch model's named modules.
+        to_numpy: A function to convert a Tensor to numpy array
 
     Returns:
         A list of Graph nodes that were built from the fx GraphModule nodes.
@@ -91,7 +93,7 @@ def nodes_builder(model: GraphModule,
             node_type = node.target
             if node_type == getattr:
                 node_has_activation = False
-                common.Logger.warning(
+                Logger.warning(
                     'Pytorch model has a parameter or constant Tensor value. This can cause unexpected behaviour when '
                     'converting the model.')
         elif node.op == PLACEHOLDER:
@@ -112,7 +114,7 @@ def nodes_builder(model: GraphModule,
             else:
                 node_type = ConstantHolder
             node_has_activation = False
-            common.Logger.warning(
+            Logger.warning(
                 'Pytorch model has a parameter or constant Tensor value. This can cause unexpected behaviour when '
                 'converting the model.')
         else:

model_compression_toolkit/core/pytorch/statistics_correction/apply_second_moment_correction.py

@@ -18,7 +18,7 @@ from typing import Any, Callable
 import torch
 from tqdm import tqdm
 
-from model_compression_toolkit import CoreConfig
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.pytorch.constants import GAMMA, BETA, MOVING_MEAN, MOVING_VARIANCE
 from model_compression_toolkit.core.pytorch.utils import set_model, to_torch_tensor

model_compression_toolkit/core/runner.py

@@ -22,7 +22,7 @@ from tqdm import tqdm
 
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import FrameworkInfo
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.fusion.layer_fusing import fusion
 from model_compression_toolkit.core.common.graph.base_graph import Graph
@@ -48,7 +48,7 @@ from model_compression_toolkit.core.common.statistics_correction.statistics_corr
 from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
 from model_compression_toolkit.core.common.substitutions.linear_collapsing_substitution import \
     linear_collapsing_substitute
-from model_compression_toolkit.core.common.target_platform.targetplatform2framework import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.visualization.final_config_visualizer import \
     WeightsFinalBitwidthConfigVisualizer, \
     ActivationFinalBitwidthConfigVisualizer
@@ -143,9 +143,9 @@ def core_runner(in_model: Any,
         weights_conf_nodes_bitwidth = tg.get_final_weights_config()
         activation_conf_nodes_bitwidth = tg.get_final_activation_config()
 
-        common.Logger.info(
+        Logger.info(
             f'Final weights bit-width configuration: {[node_b[1] for node_b in weights_conf_nodes_bitwidth]}')
-        common.Logger.info(
+        Logger.info(
             f'Final activation bit-width configuration: {[node_b[1] for node_b in activation_conf_nodes_bitwidth]}')
 
         if tb_w is not None:
@@ -259,9 +259,9 @@ def _init_tensorboard_writer(fw_info: FrameworkInfo) -> TensorboardWriter:
         A TensorBoardWriter object.
     """
     tb_w = None
-    if common.Logger.LOG_PATH is not None:
-        tb_log_dir = os.path.join(os.getcwd(), common.Logger.LOG_PATH, 'tensorboard_logs')
-        common.Logger.info(f'To use Tensorboard, please run: tensorboard --logdir {tb_log_dir}')
+    if Logger.LOG_PATH is not None:
+        tb_log_dir = os.path.join(os.getcwd(), Logger.LOG_PATH, 'tensorboard_logs')
+        Logger.info(f'To use Tensorboard, please run: tensorboard --logdir {tb_log_dir}')
         tb_w = TensorboardWriter(tb_log_dir, fw_info)
     return tb_w
 

model_compression_toolkit/exporter/__init__.py

@@ -12,3 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+
+from model_compression_toolkit.exporter.model_exporter.keras.keras_export_facade import keras_export_model, KerasExportMode
+from model_compression_toolkit.exporter.model_exporter.pytorch.pytorch_export_facade import PyTorchExportMode, pytorch_export_model
+from model_compression_toolkit.exporter.model_exporter.tflite.tflite_export_facade import tflite_export_model, TFLiteExportMode
+

model_compression_toolkit/exporter/model_exporter/__init__.py

@@ -13,6 +13,3 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.exporter.model_exporter.keras.keras_export_facade import keras_export_model, KerasExportMode
-from model_compression_toolkit.exporter.model_exporter.pytorch.pytorch_export_facade import PyTorchExportMode, pytorch_export_model
-from model_compression_toolkit.exporter.model_exporter.tflite.tflite_export_facade import tflite_export_model, TFLiteExportMode

model_compression_toolkit/exporter/model_exporter/fw_agonstic/exporter.py

@@ -17,7 +17,7 @@
 from abc import abstractmethod
 from typing import Any, Callable
 
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 
 
 class Exporter:

model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_keras_exporter.py

@@ -19,7 +19,7 @@ import keras.models
 import tensorflow as tf
 from keras.engine.base_layer import Layer
 
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.exporter.model_exporter.keras.base_keras_exporter import \
     BaseKerasExporter
 from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper

model_compression_toolkit/exporter/model_exporter/keras/keras_export_facade.py

@@ -15,8 +15,8 @@
 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
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import FOUND_TF
 
 
 class KerasExportMode(Enum):

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

@@ -15,8 +15,8 @@
 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
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import FOUND_TORCH
 
 
 class PyTorchExportMode(Enum):

model_compression_toolkit/exporter/model_exporter/tflite/fakely_quant_tflite_exporter.py

@@ -19,8 +19,8 @@ from typing import Callable
 import keras.models
 import tensorflow as tf
 
-from model_compression_toolkit import keras_load_quantized_model
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.load_model import keras_load_quantized_model
+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/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/tflite/tflite_export_facade.py

@@ -15,8 +15,8 @@
 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_TF
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import FOUND_TF
 
 
 class TFLiteExportMode(Enum):

model_compression_toolkit/exporter/model_wrapper/__init__.py

@@ -13,12 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.core.common.constants import FOUND_TF, FOUND_TORCH
+from model_compression_toolkit.exporter.model_wrapper.keras.validate_layer import is_keras_layer_exportable
+from model_compression_toolkit.exporter.model_wrapper.keras.builder.fully_quantized_model_builder import get_exportable_keras_model
 
-if FOUND_TF:
-    from model_compression_toolkit.exporter.model_wrapper.keras.validate_layer import is_keras_layer_exportable
-    from model_compression_toolkit.exporter.model_wrapper.keras.builder.fully_quantized_model_builder import get_exportable_keras_model
-
-if FOUND_TORCH:
-    from model_compression_toolkit.exporter.model_wrapper.pytorch.validate_layer import is_pytorch_layer_exportable
-    from model_compression_toolkit.exporter.model_wrapper.pytorch.builder.fully_quantized_model_builder import get_exportable_pytorch_model
+from model_compression_toolkit.exporter.model_wrapper.pytorch.validate_layer import is_pytorch_layer_exportable
+from model_compression_toolkit.exporter.model_wrapper.pytorch.builder.fully_quantized_model_builder import get_exportable_pytorch_model

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

@@ -14,46 +14,53 @@
 # ==============================================================================
 from typing import Tuple
 
-import tensorflow as tf
-from tensorflow.keras.layers import Layer
 
 from model_compression_toolkit import quantizers_infrastructure as qi
 from model_compression_toolkit.core import common
 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.core.keras.back2framework.keras_model_builder import KerasModelBuilder
-from model_compression_toolkit.exporter.model_wrapper.keras.builder.node_to_quantizers import \
-    get_quantization_quantizers
-
-
-def _get_wrapper(node: common.BaseNode,
-                 layer: Layer) -> qi.KerasQuantizationWrapper:
-    """
-    A function which takes a computational graph node and a keras layer and perform the quantization wrapping
-    Args:
-        n: A node of mct graph.
-        layer: A keras layer
-
-    Returns: Wrapped layer with weights quantizers and activation quantizers
-
-    """
-    weights_quantizers, activation_quantizers = get_quantization_quantizers(node)
-    return qi.KerasQuantizationWrapper(layer, weights_quantizers, activation_quantizers)
-
-
-def get_exportable_keras_model(graph: Graph) -> Tuple[tf.keras.models.Model,UserInformation]:
-    """
-    Convert graph to an exportable Keras model (model with all quantization parameters).
-    An exportable model can then be exported using model_exporter, to retrieve the
-    final exported model.
-
-    Args:
-        graph: Graph to convert to an exportable Keras model.
-
-    Returns:
-        Exportable Keras model and user information.
-    """
-    exportable_model, user_info = KerasModelBuilder(graph=graph,
-                                         wrapper=_get_wrapper).build_model()
-    exportable_model.trainable = False
-    return exportable_model, user_info
+from model_compression_toolkit.logger import Logger
+
+if FOUND_TF:
+    import tensorflow as tf
+    from tensorflow.keras.layers import Layer
+    from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
+    from model_compression_toolkit.exporter.model_wrapper.keras.builder.node_to_quantizers import get_quantization_quantizers
+
+    def _get_wrapper(node: common.BaseNode,
+                     layer: Layer) -> qi.KerasQuantizationWrapper:
+        """
+        A function which takes a computational graph node and a keras layer and perform the quantization wrapping
+        Args:
+            n: A node of mct graph.
+            layer: A keras layer
+
+        Returns: Wrapped layer with weights quantizers and activation quantizers
+
+        """
+        weights_quantizers, activation_quantizers = get_quantization_quantizers(node)
+        return qi.KerasQuantizationWrapper(layer, weights_quantizers, activation_quantizers)
+
+
+    def get_exportable_keras_model(graph: Graph) -> Tuple[tf.keras.models.Model, UserInformation]:
+        """
+        Convert graph to an exportable Keras model (model with all quantization parameters).
+        An exportable model can then be exported using model_exporter, to retrieve the
+        final exported model.
+
+        Args:
+            graph: Graph to convert to an exportable Keras model.
+
+        Returns:
+            Exportable Keras model and user information.
+        """
+        exportable_model, user_info = KerasModelBuilder(graph=graph,
+                                             wrapper=_get_wrapper).build_model()
+        exportable_model.trainable = False
+        return exportable_model, user_info
+else:
+    def get_exportable_keras_model(*args, **kwargs):  # pragma: no cover
+        Logger.error('Installing tensorflow and tensorflow_model_optimization is mandatory '
+                     'when using get_exportable_keras_model. '
+                     'Could not find Tensorflow package.')

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

@@ -14,16 +14,15 @@
 # ==============================================================================
 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
-from model_compression_toolkit.core.common.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
+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
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.quantizers import constants as qi_keras_consts
 
 def get_inferable_quantizer_kwargs(node: BaseNode,
                                    quantization_target: QuantizationTarget) -> Dict[str, Any]:
@@ -44,19 +43,29 @@ def get_inferable_quantizer_kwargs(node: BaseNode,
         # Return the appropriate quantization parameters based on the quantization method
         if quantization_method in [QuantizationMethod.POWER_OF_TWO,
                                    QuantizationMethod.SYMMETRIC]:
-            return {'num_bits': node_w_qc.weights_n_bits,
-                    'threshold': list(node_w_qc.weights_quantization_params[THRESHOLD].flatten()),
-                    'per_channel': node_w_qc.weights_per_channel_threshold,
-                    'channel_axis': node_w_qc.weights_channels_axis,
-                    'input_rank': len(node_w_qc.weights_quantization_params[THRESHOLD].shape)}
+            return {qi_keras_consts.NUM_BITS: node_w_qc.weights_n_bits,
+                    qi_keras_consts.THRESHOLD: list(node_w_qc.weights_quantization_params[THRESHOLD].flatten()),
+                    qi_keras_consts.PER_CHANNEL: node_w_qc.weights_per_channel_threshold,
+                    qi_keras_consts.CHANNEL_AXIS: node_w_qc.weights_channels_axis,
+                    qi_keras_consts.INPUT_RANK: len(node_w_qc.weights_quantization_params[THRESHOLD].shape)}
 
         elif quantization_method in [QuantizationMethod.UNIFORM]:
-            return {'num_bits': node_w_qc.weights_n_bits,
-                    'per_channel': node_w_qc.weights_per_channel_threshold,
-                    'min_range': list(node_w_qc.weights_quantization_params[RANGE_MIN].flatten()),
-                    'max_range': list(node_w_qc.weights_quantization_params[RANGE_MAX].flatten()),
-                    'channel_axis': node_w_qc.weights_channels_axis,
-                    'input_rank': len(node_w_qc.weights_quantization_params[RANGE_MIN].shape)}
+            return {qi_keras_consts.NUM_BITS: node_w_qc.weights_n_bits,
+                    qi_keras_consts.PER_CHANNEL: node_w_qc.weights_per_channel_threshold,
+                    qi_keras_consts.MIN_RANGE: list(node_w_qc.weights_quantization_params[RANGE_MIN].flatten()),
+                    qi_keras_consts.MAX_RANGE: list(node_w_qc.weights_quantization_params[RANGE_MAX].flatten()),
+                    qi_keras_consts.CHANNEL_AXIS: node_w_qc.weights_channels_axis,
+                    qi_keras_consts.INPUT_RANK: len(node_w_qc.weights_quantization_params[RANGE_MIN].shape)}
+
+        elif quantization_method in [QuantizationMethod.LUT_SYM_QUANTIZER, QuantizationMethod.LUT_POT_QUANTIZER]:
+            return {qi_keras_consts.NUM_BITS: node_w_qc.weights_n_bits,
+                    qi_keras_consts.PER_CHANNEL: node_w_qc.weights_per_channel_threshold,
+                    qi_keras_consts.CLUSTER_CENTERS: node_w_qc.weights_quantization_params[CLUSTER_CENTERS],
+                    qi_keras_consts.THRESHOLD: list(node_w_qc.weights_quantization_params[SCALE_PER_CHANNEL].flatten()),
+                    qi_keras_consts.CHANNEL_AXIS: node_w_qc.weights_channels_axis,
+                    # TODO: how to pass multiplier nbits and eps for a specific node?
+                    qi_keras_consts.INPUT_RANK: len(node_w_qc.weights_quantization_params[SCALE_PER_CHANNEL].shape)}
+
         else:
             Logger.critical(f'Not supported quantization method for inferable quantizers.')  # pragma: no cover
 
@@ -68,16 +77,24 @@ def get_inferable_quantizer_kwargs(node: BaseNode,
         # Return the appropriate quantization parameters based on the quantization method
         if quantization_method in [QuantizationMethod.POWER_OF_TWO,
                                    QuantizationMethod.SYMMETRIC]:
-            return {'num_bits': node_qc.activation_n_bits,
+            return {qi_keras_consts.NUM_BITS: node_qc.activation_n_bits,
                     # In activation quantization is per-tensor only - thus we hold the threshold as a list with a len of 1
-                    'threshold': [node_qc.activation_quantization_params[THRESHOLD]],
-                    'signed': node_qc.activation_quantization_params[SIGNED]}
+                    qi_keras_consts.THRESHOLD: [node_qc.activation_quantization_params[THRESHOLD]],
+                    qi_keras_consts.SIGNED: node_qc.activation_quantization_params[SIGNED]}
 
         elif quantization_method in [QuantizationMethod.UNIFORM]:
-            return {'num_bits': node_qc.activation_n_bits,
+            return {qi_keras_consts.NUM_BITS: node_qc.activation_n_bits,
                     # In activation quantization is per-tensor only - thus we hold the min/max as a list with a len of 1
-                    'min_range': [node_qc.activation_quantization_params[RANGE_MIN]],
-                    'max_range': [node_qc.activation_quantization_params[RANGE_MAX]]}
+                    qi_keras_consts.MIN_RANGE: [node_qc.activation_quantization_params[RANGE_MIN]],
+                    qi_keras_consts.MAX_RANGE: [node_qc.activation_quantization_params[RANGE_MAX]]}
+
+        elif quantization_method in [QuantizationMethod.LUT_POT_QUANTIZER]:
+            return {qi_keras_consts.NUM_BITS: node_qc.activation_n_bits,
+                    qi_keras_consts.SIGNED: node_qc.activation_quantization_params[SIGNED],
+                    qi_keras_consts.CLUSTER_CENTERS: node_qc.activation_quantization_params[CLUSTER_CENTERS],
+                    qi_keras_consts.THRESHOLD: [node_qc.activation_quantization_params[THRESHOLD]]
+                    # TODO: how to pass multiplier nbits and eps for a specific node?
+                    }
         else:
             Logger.critical(f'Not supported quantization method for inferable quantizers.')  # pragma: no cover
     else: