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

model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_search.py

@@ -18,7 +18,7 @@ from typing import Any, Tuple, Dict
 
 import numpy as np
 
-from model_compression_toolkit.core.common.constants import MIN_THRESHOLD, DEFAULT_TOL, DEFAULT_DEC_FACTOR, \
+from model_compression_toolkit.constants import MIN_THRESHOLD, DEFAULT_TOL, DEFAULT_DEC_FACTOR, \
     SYMMETRIC_TENSOR_PER_CHANNEL_N_INTERVALS, SYMMETRIC_TENSOR_PER_CHANNEL_N_ITER, SYMMETRIC_TENSOR_DEC_FREQ, \
     SYMMETRIC_TENSOR_PER_CHANNEL_DEC_FREQ, SYMMETRIC_TENSOR_N_INTERVALS, SYMMETRIC_TENSOR_N_ITER, \
     UNIFORM_TENSOR_PER_CHANNEL_N_ITER, UNIFORM_TENSOR_N_ITER, SYMMETRIC_HISTOGRAM_DEC_FREQ, SYMMETRIC_HISTOGRAM_N_ITER, \

model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_weights_computation.py

@@ -16,7 +16,7 @@ from typing import Dict, Any, Tuple
 
 import numpy as np
 
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.defaultdict import DefaultDict
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig

model_compression_toolkit/core/common/quantization/quantization_params_generation/symmetric_selection.py

@@ -15,7 +15,7 @@
 import numpy as np
 
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
-from model_compression_toolkit.core.common.constants import MIN_THRESHOLD, THRESHOLD
+from model_compression_toolkit.constants import MIN_THRESHOLD, THRESHOLD
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.error_functions import \
     get_threshold_selection_tensor_error_function, get_threshold_selection_histogram_error_function, _kl_error_histogram
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_search import \

model_compression_toolkit/core/common/quantization/quantization_params_generation/uniform_selection.py

@@ -15,7 +15,7 @@
 import numpy as np
 
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
-from model_compression_toolkit.core.common.constants import MIN_THRESHOLD, RANGE_MIN, RANGE_MAX
+from model_compression_toolkit.constants import MIN_THRESHOLD, RANGE_MIN, RANGE_MAX
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_search import \
     qparams_uniform_selection_tensor_search, qparams_uniform_selection_histogram_search
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.error_functions import \

model_compression_toolkit/core/common/quantization/quantize_graph_weights.py

@@ -20,6 +20,7 @@ from model_compression_toolkit.core.common.framework_implementation import Frame
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.quantization.quantize_node import get_quantized_kernel_by_weights_qc
+from model_compression_toolkit.logger import Logger
 
 
 def quantize_graph_weights(graph: Graph,
@@ -47,7 +48,7 @@ def quantize_graph_weights(graph: Graph,
                                                                                     n.final_weights_quantization_cfg,
                                                                                     fw_impl=fw_impl)
 
-            common.Logger.debug(
+            Logger.debug(
                 f'Node name: {n.name} has the following quantization params: '
                 f'{str(n.final_weights_quantization_cfg.weights_quantization_params)}')
 

model_compression_toolkit/core/common/quantization/quantize_node.py

@@ -15,7 +15,7 @@
 
 
 from model_compression_toolkit.core import common
-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.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
@@ -46,7 +46,7 @@ def get_quantized_kernel_by_weights_qc(fw_info: FrameworkInfo,
     # If weights should be quantized per-channel but a kernel channels mapping is missing.
     if weights_qc.weights_per_channel_threshold and fw_info.kernel_channels_mapping is \
             None:
-        common.Logger.warning(
+        Logger.warning(
             'Weights Per Channel Quantization requires channel mapping function but framework info '
             'does not contain one')
     output_channels_axis, input_channels_axis = get_channels_axis(weights_qc,

model_compression_toolkit/core/common/quantization/quantizers/kmeans_quantizer.py

@@ -16,7 +16,7 @@
 from sklearn.cluster import KMeans
 import numpy as np
 
-from model_compression_toolkit.core.common.constants import CLUSTER_CENTERS, MIN_THRESHOLD, SCALE_PER_CHANNEL
+from model_compression_toolkit.constants import CLUSTER_CENTERS, MIN_THRESHOLD, SCALE_PER_CHANNEL
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import kmeans_assign_clusters
 
 

model_compression_toolkit/core/common/quantization/quantizers/lut_kmeans_quantizer.py

@@ -15,7 +15,7 @@
 
 import numpy as np
 
-from model_compression_toolkit.core.common.constants import CLUSTER_CENTERS, SCALE_PER_CHANNEL, \
+from model_compression_toolkit.constants import CLUSTER_CENTERS, SCALE_PER_CHANNEL, \
     MULTIPLIER_N_BITS
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import kmeans_assign_clusters, \
     get_quantized_tensor, int_quantization_with_threshold

model_compression_toolkit/core/common/quantization/quantizers/quantizers_helpers.py

@@ -17,8 +17,10 @@
 from typing import Tuple, List
 import numpy as np
 
-from model_compression_toolkit.core.common.constants import MIN_THRESHOLD, EPS
+from model_compression_toolkit.constants import MIN_THRESHOLD, EPS
 from model_compression_toolkit.core import common
+from model_compression_toolkit.logger import Logger
+
 
 def max_power_of_two(x: np.ndarray,
                      min_threshold: float = MIN_THRESHOLD) -> np.ndarray:
@@ -236,7 +238,7 @@ def get_tensor_max(tensor_data: np.ndarray,
 
     """
     if n_bits < 1:
-        common.Logger.error("n_bits must be positive")
+        Logger.error("n_bits must be positive")
     if is_uniform_quantization:
         expansion_factor = 1.0
     elif n_bits == 1:

model_compression_toolkit/core/common/quantization/quantizers/uniform_quantizers.py

@@ -15,8 +15,8 @@
 
 import numpy as np
 
-from model_compression_toolkit.core.common.logger import Logger
-from model_compression_toolkit.core.common.constants import RANGE_MIN, RANGE_MAX, THRESHOLD
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import RANGE_MIN, RANGE_MAX, THRESHOLD
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import uniform_quantize_tensor, \
     quantize_tensor
 

model_compression_toolkit/core/common/quantization/set_node_quantization_config.py

@@ -17,7 +17,8 @@
 import copy
 from typing import List
 
-from model_compression_toolkit.core.common import Logger, BaseNode
+from model_compression_toolkit.core.common import BaseNode
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
@@ -72,7 +73,7 @@ def set_quantization_configs_to_node(node: BaseNode,
         tpc: TargetPlatformCapabilities to get default OpQuantizationConfig.
         mixed_precision_enable: is mixed precision enabled
     """
-    node_qc_options = tpc.get_qco_by_node(node)
+    node_qc_options = node.get_qco(tpc)
 
     # Create QC candidates for weights and activation combined
     weight_channel_axis = fw_info.kernel_channels_mapping.get(node.type)[0]

model_compression_toolkit/core/common/similarity_analyzer.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Any, Tuple
+from typing import Any
 
 import numpy as np
 
-from model_compression_toolkit.core.common.constants import EPS
+from model_compression_toolkit.constants import EPS
 
 #########################
 #  Helpful functions

model_compression_toolkit/core/common/statistics_correction/apply_bias_correction_to_graph.py

@@ -14,12 +14,12 @@
 # ==============================================================================
 import copy
 
-from model_compression_toolkit import CoreConfig
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 
 
-def apply_bias_correction_to_graph(graph: Graph,
+def apply_bias_correction_to_graph(graph_to_apply_bias_correction: Graph,
                                    core_config: CoreConfig,
                                    fw_impl: FrameworkImplementation) -> Graph:
     """
@@ -27,7 +27,7 @@ def apply_bias_correction_to_graph(graph: Graph,
     correction term in it), and apply the bias correction for each node in the graph.
 
     Args:
-        graph: Graph to apply bias correction to.
+        graph_to_apply_bias_correction: Graph to apply bias correction to.
         core_config: CoreConfig containing parameters of how the model should be quantized.
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
 
@@ -35,6 +35,7 @@ def apply_bias_correction_to_graph(graph: Graph,
         Graph with bias correction apply to it's nodes.
     """
 
+    graph = copy.deepcopy(graph_to_apply_bias_correction)
     for n in graph.nodes:
         if n.is_weights_quantization_enabled() and core_config.quantization_config.weights_bias_correction \
                 and not n.final_weights_quantization_cfg.weights_second_moment_correction:

model_compression_toolkit/core/common/statistics_correction/compute_bias_correction_of_graph.py

@@ -18,12 +18,13 @@ from typing import Any
 
 import numpy as np
 
-from model_compression_toolkit import CoreConfig
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.common import BaseNode, Logger, Graph
+from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.quantization.quantize_node import get_quantized_kernel_by_weights_qc
 from model_compression_toolkit.core.common.collectors.statistics_collector import BaseStatsCollector
+from model_compression_toolkit.logger import Logger
 
 
 def compute_bias_correction_of_graph(graph: Graph,

model_compression_toolkit/core/common/substitutions/batchnorm_reconstruction.py

@@ -20,7 +20,7 @@ from typing import Callable
 import numpy as np
 
 from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common import Logger
+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
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher

model_compression_toolkit/core/common/substitutions/batchnorm_refusing.py

@@ -23,8 +23,8 @@ from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import EdgeMatcher, NodeOperationMatcher
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
-from model_compression_toolkit.core.common.constants import THRESHOLD, RANGE_MIN, RANGE_MAX
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.constants import THRESHOLD, RANGE_MIN, RANGE_MAX
+from model_compression_toolkit.logger import Logger
 
 
 class BatchNormalizationRefusing(common.BaseSubstitution):

model_compression_toolkit/core/common/substitutions/linear_collapsing.py

@@ -17,7 +17,7 @@
 import copy
 import numpy as np
 from typing import Tuple, Callable
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import EdgeMatcher, NodeOperationMatcher

model_compression_toolkit/core/common/substitutions/shift_negative_activation.py

@@ -16,9 +16,9 @@ import copy
 import numpy as np
 from typing import List, Tuple, Any, Callable
 
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common import FrameworkInfo, Graph, BaseNode
-from model_compression_toolkit.core.common.constants import THRESHOLD, SIGNED, SHIFT_NEGATIVE_NON_LINEAR_NUM_BITS
+from model_compression_toolkit.constants import THRESHOLD, SIGNED, SHIFT_NEGATIVE_NON_LINEAR_NUM_BITS
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
 from model_compression_toolkit.core.common.quantization.set_node_quantization_config import create_node_activation_qc, \
@@ -356,7 +356,7 @@ def shift_negative_function(graph: Graph,
                     bypass_candidate_qc.activation_quantization_cfg.activation_quantization_params[SIGNED] = False
                     graph.shift_stats_collector(bypass_node, np.array(shift_value))
 
-    add_node_qco = graph.tpc.get_qco_by_node(add_node).quantization_config_list
+    add_node_qco = add_node.get_qco(graph.tpc).quantization_config_list
     for op_qc_idx, candidate_qc in enumerate(add_node.candidates_quantization_cfg):
         candidate_qc.weights_quantization_cfg.enable_weights_quantization = False
 
@@ -495,7 +495,7 @@ def apply_shift_negative_correction(graph: Graph,
     nodes = list(graph.nodes())
     for n in nodes:
         # Skip substitution if QuantizationMethod is uniform.
-        node_qco = graph.tpc.get_qco_by_node(n)
+        node_qco = n.get_qco(graph.tpc)
         if any([op_qc.activation_quantization_method is QuantizationMethod.UNIFORM
                 for op_qc in node_qco.quantization_config_list]):
             continue

model_compression_toolkit/core/common/substitutions/virtual_activation_weights_composition.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 from model_compression_toolkit.core.common import BaseNode, Graph, BaseSubstitution
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.graph.virtual_activation_weights_node import VirtualActivationWeightsNode
 
 

model_compression_toolkit/core/common/substitutions/weights_activation_split.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 import itertools
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common import BaseNode, Graph, BaseSubstitution
 from model_compression_toolkit.core.common.graph.virtual_activation_weights_node import VirtualSplitWeightsNode, \
     VirtualSplitActivationNode

model_compression_toolkit/core/common/visualization/tensorboard_writer.py

@@ -31,7 +31,7 @@ from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto
 from tensorboard.summary.writer.event_file_writer import EventFileWriter
 from typing import List, Any, Dict
 from networkx import topological_sort
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.collectors.statistics_collector import BaseStatsCollector
 

model_compression_toolkit/core/keras/back2framework/factory_model_builder.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
 from model_compression_toolkit.core.keras.back2framework.float_model_builder import FloatKerasModelBuilder
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder

model_compression_toolkit/core/keras/back2framework/float_model_builder.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from typing import List
 
-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.keras.back2framework.keras_model_builder import KerasModelBuilder
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO

model_compression_toolkit/core/keras/back2framework/keras_model_builder.py

@@ -13,16 +13,15 @@
 # limitations under the License.
 # ==============================================================================
 
-from abc import abstractmethod
-
 import tensorflow as tf
 from keras.engine.input_layer import InputLayer
 from keras.models import Model, clone_model
 from packaging import version
 
+from model_compression_toolkit.constants import INPUT_BASE_NAME
 from model_compression_toolkit.core.common.back2framework.base_model_builder import BaseModelBuilder
 from model_compression_toolkit.core.common.user_info import UserInformation
-from model_compression_toolkit.core.common.constants import INPUT_BASE_NAME
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.activation_quantization_holder import ActivationQuantizationHolder
 
 # As from Tensorflow 2.6, keras is a separate package and some classes should be imported differently.
 if version.parse(tf.__version__) < version.parse("2.6"):
@@ -38,7 +37,6 @@ else:
 from typing import Any, Dict, List, Tuple, Callable
 from tensorflow.python.util.object_identity import Reference as TFReference
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
-from model_compression_toolkit.core.common.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
@@ -95,7 +93,8 @@ class KerasModelBuilder(BaseModelBuilder):
                  append2output=None,
                  fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                  return_float_outputs: bool = False,
-                 wrapper: Callable = None):
+                 wrapper: Callable = None,
+                 get_activation_quantizer_holder_fn: Callable=None):
         """
 
         Args:
@@ -104,6 +103,8 @@ class KerasModelBuilder(BaseModelBuilder):
             fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
             wrapper: A function wrapper keras Layers.
+            get_activation_quantizer_holder_fn: Function to retrieve a quantization holder for a node.
+
         """
 
         super().__init__(graph,
@@ -114,6 +115,19 @@ class KerasModelBuilder(BaseModelBuilder):
         # Build an OperationHandler to handle conversions from graph nodes to Keras operators.
         self.oh = OperationHandler(self.graph)
         self.wrapper = wrapper
+        self.get_activation_quantizer_holder = get_activation_quantizer_holder_fn
+
+    @property
+    def use_activation_holder_during_model_building(self) -> bool:
+        """
+
+        Returns: Whether the model builder uses ActivationQuantizationHolder during
+        model building (by adding it as a layer when converting the graph to the Keras model)
+        or not. If so - the model builder expects the activation quantizers to not be wrapped
+        in KerasQuantizeWrapper that was received in its init.
+
+        """
+        return self.get_activation_quantizer_holder is not None
 
     def _quantize_node_activations(self,
                                    node: BaseNode,
@@ -187,9 +201,8 @@ class KerasModelBuilder(BaseModelBuilder):
                 node_to_output_tensors_dict.update({n: [out_tensors_of_n]})
                 node_to_output_tensors_dict_float.update({n: [out_tensors_of_n_float]})
 
-        # convert node_to_output_tensors_dict keys to nodes' names since oh.node_sort contains different objects
-        # than
-        # original graph nodes.
+        # convert node_to_output_tensors_dict keys to nodes' names since oh.node_sort
+        # contains different objects than original graph nodes.
         node_name_to_outtensors = self._convert_node2name(node_to_output_tensors_dict)
         node_name_to_outtensors_float = self._convert_node2name(node_to_output_tensors_dict_float)
 
@@ -214,9 +227,12 @@ class KerasModelBuilder(BaseModelBuilder):
             def _wrap(layer):
                 _node = self.oh.layer_to_node_dict.get(layer)
                 if _node is not None:
-                    return self.wrapper(_node, layer)
-                elif is_layer_fake_quant(layer):
+                    return self.wrapper(_node,
+                                        layer)
+
+                elif is_layer_fake_quant(layer) or isinstance(layer, ActivationQuantizationHolder):
                     return layer
+
                 raise Exception(  # pragma: no cover
                     f'Mismatch between keras model and graph cant find node named: '
                     f'{get_node_name_from_layer(layer)}')
@@ -278,13 +294,9 @@ class KerasModelBuilder(BaseModelBuilder):
         """
         if len(input_tensors) == 0:  # Placeholder handling
             out_tensors_of_n_float = input_nodes_to_input_tensors[n]
-            if self.wrapper is not None:
-                # if a wrapper is defined, add an identity layer for cloning. The Identity will be warpped
-                out_tensors_of_n = op_func(out_tensors_of_n_float)
-            elif n.is_activation_quantization_enabled():
-                out_tensors_of_n = self._quantize_node_activations(n, out_tensors_of_n_float)
-            else:
-                out_tensors_of_n = out_tensors_of_n_float
+            out_tensors_of_n = self._run_operation_activation_quantization(n,
+                                                                           out_tensors_of_n_float,
+                                                                           op_func)
         else:
             input_tensors = [tensor for tensor_list in input_tensors for tensor in tensor_list]  # flat list of lists
             # Build a functional node using its args
@@ -299,11 +311,9 @@ class KerasModelBuilder(BaseModelBuilder):
                 if len(input_tensors) == 1:
                     input_tensors = input_tensors[0]
                 out_tensors_of_n_float = op_func(input_tensors)
-            out_tensors_of_n = out_tensors_of_n_float
 
-            # Add a fake quant node if the node has an activation threshold and a wrapper isn't defined
-            if n.is_activation_quantization_enabled() and self.wrapper is None:
-                out_tensors_of_n = self._quantize_node_activations(n, out_tensors_of_n_float)
+            out_tensors_of_n = self._run_operation_activation_quantization(n,
+                                                                           out_tensors_of_n_float)
 
         # Save a mapping from the layer that created the tensor to the node (as this layer is not the
         # same instance as op_func. We do this to solve an issue that names are different between these
@@ -318,3 +328,38 @@ class KerasModelBuilder(BaseModelBuilder):
             self.oh.layer_to_node_dict[layer_from_tensor] = n
 
         return out_tensors_of_n, out_tensors_of_n_float
+
+    def _run_operation_activation_quantization(self,
+                                               node: BaseNode,
+                                               node_outputs: List[TFReference],
+                                               identity_layer: Layer = None):
+        """
+        Quantize node's activations
+
+        Args:
+            node: Node to quantize its activations
+            node_outputs: Output tensors of the float node.
+            identity_layer: Identity layer (should be passed only when quantizing input layers)
+
+        Returns:
+            Quantized node's outputs.
+        """
+        if self.wrapper is not None:
+            # If identity layer was passed, use it for inference
+            # (needed since wrapping an Input layer can not be wrapped)
+            if identity_layer is not None:
+                node_outputs = identity_layer(node_outputs)
+
+            # In case the activation quantizer is attached out of the wrapper we use get_activation_quantizer_holder
+            # for the activation quantization holder (if the node's activations are quantized)
+            if node.is_activation_quantization_enabled() and self.use_activation_holder_during_model_building:
+                activation_quantizer_holder = self.get_activation_quantizer_holder(node)
+                quantized_node_outputs = activation_quantizer_holder(node_outputs)
+                return quantized_node_outputs
+
+        elif node.is_activation_quantization_enabled():  # Used only when old exporter is used
+            quantized_node_outputs = self._quantize_node_activations(node,
+                                                                     node_outputs)
+            return quantized_node_outputs
+
+        return node_outputs

model_compression_toolkit/core/keras/back2framework/mixed_precision_model_builder.py

@@ -36,7 +36,7 @@ else:
     from keras.layers.core import TFOpLambda, SlicingOpLambda
 
 from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapper import QuantizeWrapper
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO

model_compression_toolkit/core/keras/back2framework/model_gradients.py

@@ -26,13 +26,13 @@ else:
 
 from typing import Any, Dict, List, Tuple
 from tensorflow.python.util.object_identity import Reference as TFReference
-from model_compression_toolkit.core.common.constants import EPS, MIN_JACOBIANS_ITER, JACOBIANS_COMP_TOLERANCE
+from model_compression_toolkit.constants import EPS, MIN_JACOBIANS_ITER, JACOBIANS_COMP_TOLERANCE
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX
 from model_compression_toolkit.core.keras.back2framework.instance_builder import OperationHandler
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 
 
 def build_input_tensors_list(node: BaseNode,

model_compression_toolkit/core/keras/back2framework/quantized_model_builder.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from typing import List
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder

model_compression_toolkit/core/keras/constants.py

@@ -99,13 +99,6 @@ OUTPUT_BIAS = '/attention_output/bias'
 # ReLU bound constants
 RELU_POT_BOUND = 8.0
 
-# Supported TP models names for Tensorflow:
-DEFAULT_TP_MODEL = 'default'
-IMX500_TP_MODEL = 'imx500'
-TFLITE_TP_MODEL = 'tflite'
-QNNPACK_TP_MODEL = 'qnnpack'
-
-
 # TFOpLambda functions:
 ADD = 'add'
 PAD = 'pad'

model_compression_toolkit/core/keras/default_framework_info.py

@@ -25,9 +25,9 @@ else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU
 
 from model_compression_toolkit.core.common.defaultdict import DefaultDict
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo, ChannelAxis
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
-from model_compression_toolkit.core.common.constants import SOFTMAX_THRESHOLD
+from model_compression_toolkit.constants import SOFTMAX_THRESHOLD
 from model_compression_toolkit.core.keras.constants import SOFTMAX, LINEAR, RELU, SWISH, SIGMOID, IDENTITY, TANH, SELU, \
     KERNEL, DEPTHWISE_KERNEL
 from model_compression_toolkit.core.keras.quantizer.fake_quant_builder import power_of_two_quantization, symmetric_quantization, uniform_quantization

model_compression_toolkit/core/keras/graph_substitutions/substitutions/activation_decomposition.py

@@ -17,7 +17,7 @@
 from tensorflow.keras.layers import Dense, DepthwiseConv2D, Conv2D, Conv2DTranspose, Activation, SeparableConv2D
 
 from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common.constants import FLOAT_32, DATA_TYPE
+from model_compression_toolkit.constants import FLOAT_32, DATA_TYPE
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, \
     NodeFrameworkAttrMatcher

model_compression_toolkit/core/keras/graph_substitutions/substitutions/input_scaling.py

@@ -23,7 +23,7 @@ from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, EdgeMatcher, WalkMatcher
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
-from model_compression_toolkit.core.common.constants import THRESHOLD
+from model_compression_toolkit.constants import THRESHOLD
 from model_compression_toolkit.core.keras.constants import KERNEL
 
 input_node = NodeOperationMatcher(InputLayer)

model_compression_toolkit/core/keras/graph_substitutions/substitutions/linear_collapsing.py

@@ -21,7 +21,7 @@ from model_compression_toolkit.core.common.graph.graph_matchers import NodeOpera
 from model_compression_toolkit.core.common.substitutions.linear_collapsing import Conv2DCollapsing
 from model_compression_toolkit.core.keras.constants import KERNEL, KERNEL_SIZE, STRIDES, DILATIONS, LINEAR, \
     ACTIVATION, BIAS, USE_BIAS, LAYER_NAME, FILTERS, PADDING, GROUPS, DATA_FORMAT
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 
 
 def linear_collapsing_node_matchers() -> Tuple[NodeOperationMatcher, NodeOperationMatcher]:

model_compression_toolkit/core/keras/graph_substitutions/substitutions/multi_head_attention_decomposition.py

@@ -23,17 +23,16 @@ else:
     from keras.layers.core import TFOpLambda
     from keras.layers import MultiHeadAttention, Conv2D, Softmax, Concatenate, Reshape, Permute
 
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.graph.base_graph import Graph, BaseNode, OutTensor
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
-from model_compression_toolkit.core.common.constants import REUSE, REUSE_GROUP
-from model_compression_toolkit.core.keras.reader.node_builder import REUSED_IDENTIFIER
+from model_compression_toolkit.constants import REUSE, REUSE_GROUP
 from model_compression_toolkit.core.keras.constants import KERNEL, BIAS, USE_BIAS, NUM_HEADS, KEY_DIM, VALUE_DIM, \
     QUERY_SHAPE, KEY_SHAPE, VALUE_SHAPE, OUTPUT_SHAPE, ATTENTION_AXES, ACTIVATION, LINEAR, FILTERS, \
     FUNCTION, DIMS, TARGET_SHAPE, F_STRIDED_SLICE, F_STACK, Q_KERNEL, Q_BIAS, K_KERNEL, K_BIAS, V_KERNEL, V_BIAS, \
-    OUTPUT_KERNEL, OUTPUT_BIAS, F_MATMUL, TRANSPOSE_B, KERNEL_SIZE, AXIS, F_STRIDED_SLICE_BEGIN, F_STRIDED_SLICE_END
+    OUTPUT_KERNEL, OUTPUT_BIAS, F_MATMUL, KERNEL_SIZE, AXIS, F_STRIDED_SLICE_BEGIN, F_STRIDED_SLICE_END
 
 
 class MHAParams:

model_compression_toolkit/core/keras/graph_substitutions/substitutions/relu_bound_to_power_of_2.py

@@ -23,6 +23,7 @@ from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, WalkMatcher
 from model_compression_toolkit.core.keras.constants import KERNEL, BIAS, RELU_MAX_VALUE, RELU_POT_BOUND
+from model_compression_toolkit.logger import Logger
 
 
 class ReLUBoundToPowerOfTwo(common.BaseSubstitution):
@@ -81,7 +82,7 @@ class ReLUBoundToPowerOfTwo(common.BaseSubstitution):
         scale_factor = max_value / self.threshold
 
         non_linear_node.framework_attr[RELU_MAX_VALUE] = np.float32(self.threshold)
-        common.Logger.debug(
+        Logger.debug(
             f"Node named:{non_linear_node.name} max value change "
             f"to:{non_linear_node.framework_attr[RELU_MAX_VALUE]}")