mct-nightly 1.7.1.31122022.post351__py3-none-any.whl → 1.8.0.1042023.post423__py3-none-any.whl

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

@@ -60,7 +60,8 @@ def lut_kmeans_tensor(tensor_data: np.ndarray,
         the thresholds per channel and the multiplier num bits.
     """
     if n_bits > MULTIPLIER_N_BITS:
-        Logger.critical(f'Look-Up-Table bit configuration has {n_bits} bits, but must be less or equal to {MULTIPLIER_N_BITS}')
+        Logger.critical(f'Look-Up-Table bit configuration has {n_bits} bits, but must be less or equal to '
+                        f'{MULTIPLIER_N_BITS}')  # pragma: no cover
     # TODO: need to set this externally
     if len(np.unique(tensor_data.flatten())) < 2 ** n_bits:
         n_clusters = len(np.unique(tensor_data.flatten()))
@@ -115,7 +116,8 @@ def lut_kmeans_histogram(bins: np.ndarray,
     """
 
     if n_bits >= MULTIPLIER_N_BITS:
-        Logger.critical(f'Look-Up-Table bit configuration has {n_bits} bits. It must be less then {MULTIPLIER_N_BITS}')
+        Logger.critical(f'Look-Up-Table bit configuration has {n_bits} bits. It must be less then '
+                        f'{MULTIPLIER_N_BITS}')  # pragma: no cover
 
     bins_with_values = np.abs(bins)[1:][counts > 0]
     if len(np.unique(bins_with_values.flatten())) < 2 ** n_bits:

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

@@ -49,25 +49,22 @@ def get_activations_qparams(activation_quant_cfg: NodeActivationQuantizationConf
                                                                     bins_counts)
     min_value, max_value = out_stats_container.get_min_max_values()
 
-    if nodes_prior_info is not None:
-        if nodes_prior_info.is_output_bounded():
-            signed = min_value < 0
-        else:
-            signed = np.any(bins_values[:-1][bins_counts > 0] < 0)
-
-        if nodes_prior_info.is_output_bounded():
-            if activation_quant_cfg.activation_quantization_method == QuantizationMethod.POWER_OF_TWO:
-                activation_quant_cfg.activation_quantization_params_fn = \
-                    quantization_params_generation.power_of_two_no_clipping_selection_min_max
-            elif activation_quant_cfg.activation_quantization_method == QuantizationMethod.SYMMETRIC:
-                activation_quant_cfg.activation_quantization_params_fn = \
-                    quantization_params_generation.symmetric_no_clipping_selection_min_max
-            elif activation_quant_cfg.activation_quantization_method == QuantizationMethod.UNIFORM:
-                activation_quant_cfg.activation_quantization_params_fn = \
-                    quantization_params_generation.uniform_no_clipping_selection_min_max
+    if nodes_prior_info.is_output_bounded():
+        signed = min_value < 0
     else:
         signed = np.any(bins_values[:-1][bins_counts > 0] < 0)
 
+    if nodes_prior_info.is_output_bounded():
+        if activation_quant_cfg.activation_quantization_method == QuantizationMethod.POWER_OF_TWO:
+            activation_quant_cfg.activation_quantization_params_fn = \
+                quantization_params_generation.power_of_two_no_clipping_selection_min_max
+        elif activation_quant_cfg.activation_quantization_method == QuantizationMethod.SYMMETRIC:
+            activation_quant_cfg.activation_quantization_params_fn = \
+                quantization_params_generation.symmetric_no_clipping_selection_min_max
+        elif activation_quant_cfg.activation_quantization_method == QuantizationMethod.UNIFORM:
+            activation_quant_cfg.activation_quantization_params_fn = \
+                quantization_params_generation.uniform_no_clipping_selection_min_max
+
     activation_params = activation_quant_cfg.activation_quantization_params_fn(bins_values,
                                                                                bins_counts,
                                                                                activation_quant_cfg.l_p_value,
@@ -78,4 +75,4 @@ def get_activations_qparams(activation_quant_cfg: NodeActivationQuantizationConf
                                                                                quant_error_method=activation_quant_cfg.activation_error_method)
     activation_params.update({SIGNED: signed})
 
-    return activation_params
+    return activation_params

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

@@ -18,7 +18,7 @@ from typing import Tuple, List
 import numpy as np
 
 from model_compression_toolkit.core.common.constants import MIN_THRESHOLD, EPS
-
+from model_compression_toolkit.core import common
 
 def max_power_of_two(x: np.ndarray,
                      min_threshold: float = MIN_THRESHOLD) -> np.ndarray:
@@ -235,7 +235,14 @@ def get_tensor_max(tensor_data: np.ndarray,
     Returns: maximal value (or values).
 
     """
-    expansion_factor = 1.0 if is_uniform_quantization else np.power(2.0, n_bits - 1) / (np.power(2.0, n_bits - 1) - 1)
+    if n_bits < 1:
+        common.Logger.error("n_bits must be positive")
+    if is_uniform_quantization:
+        expansion_factor = 1.0
+    elif n_bits == 1:
+        expansion_factor = 0.0
+    else:
+        expansion_factor = np.power(2.0, n_bits - 1) / (np.power(2.0, n_bits - 1) - 1)
     if per_channel:
         output_shape = get_output_shape(tensor_data.shape, channel_axis)
         reshaped_tensor_data = reshape_tensor_for_per_channel_search(tensor_data, channel_axis)

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

@@ -15,6 +15,7 @@
 
 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.core.common.quantization.quantizers.quantizers_helpers import uniform_quantize_tensor, \
     quantize_tensor
@@ -51,9 +52,9 @@ def power_of_two_quantizer(tensor_data: np.ndarray,
     """
     threshold = quantization_params.get(THRESHOLD)
     if threshold is None:
-        raise Exception(f"{THRESHOLD} parameter must be defined in 'quantization_params'")
+        Logger.error(f"{THRESHOLD} parameter must be defined in 'quantization_params'")  # pragma: no cover
     if not threshold_is_power_of_two(threshold, per_channel):
-        raise Exception(f"Expects {THRESHOLD} parameter to be a power of two, but got {threshold}")
+        Logger.error(f"Expects {THRESHOLD} parameter to be a power of two, but got {threshold}")  # pragma: no cover
 
     return quantize_tensor(tensor_data,
                            threshold,
@@ -84,7 +85,7 @@ def symmetric_quantizer(tensor_data: np.ndarray,
     """
     threshold = quantization_params.get(THRESHOLD)
     if threshold is None:
-        raise Exception(f"{THRESHOLD} parameter must be defined in 'quantization_params'")
+        Logger.error(f"{THRESHOLD} parameter must be defined in 'quantization_params'")  # pragma: no cover
 
     return quantize_tensor(tensor_data,
                            threshold,
@@ -115,6 +116,6 @@ def uniform_quantizer(tensor_data: np.ndarray,
     range_min = quantization_params.get(RANGE_MIN)
     range_max = quantization_params.get(RANGE_MAX)
     if range_min is None or range_max is None:
-        raise Exception("'quantization range' parameters must be defined in 'quantization_params'")
+        Logger.error("'quantization range' parameters must be defined in 'quantization_params'")  # pragma: no cover
 
     return uniform_quantize_tensor(tensor_data, range_min, range_max, n_bits)

model_compression_toolkit/core/common/quantization/set_node_quantization_config.py

@@ -108,7 +108,7 @@ def create_node_activation_qc(qc: QuantizationConfig,
 
     activation_quantization_fn = fw_info.activation_quantizer_mapping.get(op_cfg.activation_quantization_method)
     if activation_quantization_fn is None:
-        Logger.critical('Unknown quantization method for activations')
+        Logger.critical('Unknown quantization method for activations')  # pragma: no cover
 
     activation_quantization_params_fn = get_activation_quantization_params_fn(op_cfg.activation_quantization_method)
 
@@ -142,14 +142,14 @@ def create_node_qc_candidate(qc: QuantizationConfig,
     weights_quantization_fn = get_weights_quantization_fn(op_cfg.weights_quantization_method)
 
     if weights_quantization_fn is None:
-        Logger.critical('Unknown quantization method for weights')
+        Logger.critical('Unknown quantization method for weights')  # pragma: no cover
 
     weights_quantization_params_fn = get_weights_quantization_params_fn(op_cfg.weights_quantization_method)
 
     # get attributes for activation quantization
     activation_quantization_fn = fw_info.activation_quantizer_mapping.get(op_cfg.activation_quantization_method)
     if activation_quantization_fn is None:
-        Logger.critical('Unknown quantization method for activations')
+        Logger.critical('Unknown quantization method for activations')  # pragma: no cover
 
     activation_quantization_params_fn = get_activation_quantization_params_fn(op_cfg.activation_quantization_method)
 

model_compression_toolkit/core/common/substitutions/batchnorm_reconstruction.py

@@ -77,6 +77,13 @@ class BatchNormalizationReconstruction(common.BaseSubstitution):
         num_nodes_before_substitution = len(graph.nodes)
         num_edges_before_substitution = len(graph.edges)
 
+        # If the linear operator is part of a reused group (it is the "base" node, or a reused node),
+        # we should skip the substitution.
+        if source_node.reuse or source_node.reuse_group is not None:
+            for qc in source_node.candidates_quantization_cfg:
+                qc.weights_quantization_cfg.weights_second_moment_correction = False
+            return graph
+
         # We apply only on nodes with folded BatchNormalization.
         if source_node.prior_info.std_output is None or source_node.prior_info.mean_output is None:
             for qc in source_node.candidates_quantization_cfg:

model_compression_toolkit/core/common/substitutions/batchnorm_refusing.py

@@ -24,6 +24,7 @@ from model_compression_toolkit.core.common.graph.graph_matchers import EdgeMatch
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 from model_compression_toolkit.core.common.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
 
 
 class BatchNormalizationRefusing(common.BaseSubstitution):
@@ -95,15 +96,22 @@ class BatchNormalizationRefusing(common.BaseSubstitution):
 
         source_node = edge_nodes[0]
 
+        # We apply only on nodes with reconstructed BatchNormalization.
+        if not source_node.final_weights_quantization_cfg.weights_second_moment_correction:
+            return graph
+
         # If the linear operator is part of a reused group (it is the "base" node, or a reused node),
         # we should skip the substitution.
         if source_node.reuse or source_node.reuse_group is not None:
-            return graph
+            Logger.exception("If the linear operator is part of a reused group we should skip the the BN folding "
+                             "substitution and SMC feature")  # pragma: no cover
 
         bn_node = edge_nodes[1]
 
         if len(graph.get_next_nodes(source_node)) > 1 or len(graph.get_prev_nodes(bn_node)) > 1:
-            return graph
+            Logger.exception(
+                "If the linear operator has multiple outputs or the bn layer has multiple inputs we should "
+                "skip the the BN folding substitution and SMC feature")  # pragma: no cover
 
         kernel = source_node.get_weights_by_keys(self.kernel_str)
         bias = source_node.get_weights_by_keys(self.bias_str)
@@ -113,9 +121,6 @@ class BatchNormalizationRefusing(common.BaseSubstitution):
         moving_variance = bn_node.get_weights_by_keys(self.moving_variance_str)
         eps = bn_node.framework_attr[self.epsilon_str]
 
-        if bias is None:
-            bias = 0.0
-
         weights_scale = gamma / np.sqrt(moving_variance + eps)
         bias = beta + (bias - moving_mean) * weights_scale
 
@@ -177,7 +182,7 @@ class BatchNormalizationRefusing(common.BaseSubstitution):
             corr_dict[THRESHOLD] = corr_threshold
             conv_bn.final_weights_quantization_cfg.set_weights_quantization_param(corr_dict)
 
-        # In case of SYMMETRIC weight quantization method, we update the range_min, range_max by weights_scale
+        # In case of UNIFORM weight quantization method, we update the range_min, range_max by weights_scale
         elif conv_bn.final_weights_quantization_cfg.weights_quantization_method == QuantizationMethod.UNIFORM:
             corr_dict = copy.deepcopy(conv_bn.final_weights_quantization_cfg.weights_quantization_params)
             original_range_min = conv_bn.final_weights_quantization_cfg.weights_quantization_params[RANGE_MIN]
@@ -189,5 +194,5 @@ class BatchNormalizationRefusing(common.BaseSubstitution):
             conv_bn.final_weights_quantization_cfg.set_weights_quantization_param(corr_dict)
 
         else:
-            raise Exception("Second moment statistics correction feature disabled for models with weights "
-                            "quantization method of Power of 2")
+            Logger.exception("Second moment statistics correction feature disabled for models with weights "
+                             "quantization method of Power of 2")  # pragma: no cover

model_compression_toolkit/core/common/substitutions/shift_negative_activation.py

@@ -16,6 +16,7 @@ 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.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.core.common.graph.graph_matchers import NodeOperationMatcher
@@ -25,7 +26,8 @@ from model_compression_toolkit.core.common.quantization.set_node_quantization_co
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_activations_computation \
     import get_activations_qparams
-from model_compression_toolkit.core.common.quantization.quantization_params_generation.error_functions import _mse_error_histogram
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.error_functions import \
+    _mse_error_histogram
 from model_compression_toolkit.core.common.quantization.quantization_params_generation import z_score_filter
 
 """
@@ -73,12 +75,12 @@ def op2d_bias_correction(op2d_node: BaseNode,
 
         # special case of depthwise_conv2d in tensorflow, where we have a depth multiplier for the filters
         if output_channel_index == input_channel_index:
-            axis_not_output_channel.remove(3) # 3 is the depth multiplier index
+            axis_not_output_channel.remove(3)  # 3 is the depth multiplier index
 
         bias_correction = shift_to_correct * np.sum(kernel, axis=tuple(axis_not_output_channel))
         op2d_node.set_weights_by_keys(bias_str, bias - bias_correction.flatten())
     else:
-        raise NotImplementedError
+        raise NotImplementedError  # pragma: no cover
 
 
 def insert_node_between_two_nodes(graph: Graph,
@@ -123,7 +125,7 @@ def insert_node_after_node(graph: Graph,
 
     last_nodes = graph.get_next_nodes(first_node)
     if len(last_nodes) != 1:
-        raise Exception('Can only insert if there is only one input')
+        Logger.error('Can only insert if there is only one input')  # pragma: no cover
     last_node = last_nodes[0]
     insert_node_between_two_nodes(graph, node_to_insert, first_node, last_node)
 
@@ -145,7 +147,7 @@ def insert_node_before_node(graph: Graph,
     """
     first_nodes = graph.get_prev_nodes(last_node)
     if len(first_nodes) != 1:
-        raise Exception('Can only insert if there is only one input')
+        Logger.error('Can only insert if there is only one input')  # pragma: no cover
     first_node = first_nodes[0]
     insert_node_between_two_nodes(graph, node_to_insert, first_node, last_node)
 
@@ -222,8 +224,8 @@ def shift_negative_function(graph: Graph,
     min_to_correct, max_value2compare = graph.get_out_stats_collector(non_linear_node).get_min_max_values()
 
     if not non_linear_node.is_all_activation_candidates_equal():
-        raise Exception("Shift negative correction is not supported for more than one activation quantization "
-                        "configuration candidate")
+        Logger.error("Shift negative correction is not supported for more than one activation quantization "
+                     "configuration candidate")  # pragma: no cover
 
     # all candidates have same activation config, so taking the first candidate for calculations
     non_linear_node_cfg_candidate = non_linear_node.candidates_quantization_cfg[0].activation_quantization_cfg
@@ -241,7 +243,8 @@ def shift_negative_function(graph: Graph,
     # taking the minimal quantized point that is still positive.
     num_q_points = 2 ** non_linear_node_cfg_candidate.activation_n_bits
     lsb = activation_threshold / num_q_points
-    q_points = np.linspace(0, activation_threshold - lsb, num_q_points).astype('float32')  # Change to type float32 to support tensorflow dtypes
+    q_points = np.linspace(0, activation_threshold - lsb, num_q_points).astype(
+        'float32')  # Change to type float32 to support tensorflow dtypes
 
     delta = q_points + min_to_correct
     delta[delta < 0] = np.inf
@@ -253,14 +256,16 @@ def shift_negative_function(graph: Graph,
                                     hist_bins, hist_count)
 
         min_mse, _th, _shift = np.inf, None, None
-        for _activation_threshold in [activation_threshold, 2*activation_threshold]:
+        for _activation_threshold in [activation_threshold, 2 * activation_threshold]:
             qparams = {THRESHOLD: _activation_threshold, SIGNED: False}
             _lsb = _activation_threshold / num_q_points
-            _q_points = np.linspace(0, _activation_threshold - _lsb, num_q_points).astype('float32')  # Change to type float32 to support tensorflow dtypes
+            _q_points = np.linspace(0, _activation_threshold - _lsb, num_q_points).astype(
+                'float32')  # Change to type float32 to support tensorflow dtypes
             for _shift_value in _q_points:
                 _hist_bins = hist_bins.astype(np.float32) + _shift_value
-                q_bins = non_linear_node_cfg_candidate.activation_quantization_fn(non_linear_node_cfg_candidate.activation_n_bits,
-                                                                                  qparams)(_hist_bins)
+                q_bins = non_linear_node_cfg_candidate.activation_quantization_fn(
+                    non_linear_node_cfg_candidate.activation_n_bits,
+                    qparams)(_hist_bins)
                 mse = _mse_error_histogram(q_bins, None, _hist_bins, hist_count)
                 if mse < min_mse:
                     min_mse = mse

model_compression_toolkit/core/common/substitutions/weights_activation_split.py

@@ -61,7 +61,7 @@ class BaseWeightsActivationSplit(BaseSubstitution):
                 # Node is not composite, therefore, can't be split
                 Logger.critical(f"The graph contains a node {node.name} with non composite candidates."
                                 f"In order to run mixed-precision search with BOPS target KPI, "
-                                f"all model layers should be composite.")
+                                f"all model layers should be composite.")  # pragma: no cover
 
         weights_node = VirtualSplitWeightsNode(node)
         activation_node = VirtualSplitActivationNode(node, self.activation_layer_type, self.fw_attr)

model_compression_toolkit/core/common/target_platform/current_tp_model.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from model_compression_toolkit.core.common.logger import Logger
+
 def get_current_tp_model():
     """
 
@@ -38,7 +40,7 @@ class CurrentTPModel:
 
         """
         if self.tp_model is None:
-            raise Exception('Target platform model is not initialized.')
+            Logger.error('Target platform model is not initialized.')  # pragma: no cover
         return self.tp_model
 
     def reset(self):

model_compression_toolkit/core/common/target_platform/targetplatform2framework/attribute_filter.py

@@ -16,6 +16,8 @@
 import operator
 from typing import Any, Callable, Dict
 
+from model_compression_toolkit.core.common.logger import Logger
+
 
 class Filter:
     """
@@ -31,7 +33,7 @@ class Filter:
         Returns:
             Whether the passed configuration matches the filter or not.
         """
-        raise Exception('Filter did not implement match')
+        raise NotImplemented('Filter did not implement match')  # pragma: no cover
 
 
 class AttributeFilter(Filter):
@@ -85,7 +87,7 @@ class AttributeFilter(Filter):
         """
 
         if not isinstance(other, AttributeFilter):
-            raise Exception("Not an attribute filter. Can not run an OR operation.")
+            Logger.error("Not an attribute filter. Can not run an OR operation.")  # pragma: no cover
         return OrAttributeFilter(self, other)
 
     def __and__(self, other: Any):
@@ -99,7 +101,7 @@ class AttributeFilter(Filter):
             AndAttributeFilter that filters with AND between the current AttributeFilter and the passed AttributeFilter.
         """
         if not isinstance(other, AttributeFilter):
-            raise Exception("Not an attribute filter. Can not run an AND operation.")
+            Logger.error("Not an attribute filter. Can not run an AND operation.")  # pragma: no cover
         return AndAttributeFilter(self, other)
 
     def match(self,
@@ -123,7 +125,7 @@ class AttributeFilter(Filter):
         Returns: A string representation for the filter.
 
         """
-        raise Exception("Filter must implement op_as_str ")
+        raise NotImplemented("Filter must implement op_as_str ")  # pragma: no cover
 
     def __repr__(self):
         return f'{self.attr} {self.op_as_str()} {self.value}'
@@ -267,3 +269,14 @@ class Eq(AttributeFilter):
         super().__init__(attr=attr, value=value, op=operator.eq)
 
     def op_as_str(self): return "="
+
+
+class Contains(AttributeFilter):
+    """
+    Filter configurations such that it matches configurations that have an attribute with a value that contains the value that Contains holds.
+    """
+
+    def __init__(self, attr: str, value: Any):
+        super().__init__(attr=attr, value=value, op=operator.contains)
+
+    def op_as_str(self): return " in "

model_compression_toolkit/core/common/target_platform/targetplatform2framework/operations_to_layers.py

@@ -131,9 +131,7 @@ class OperationsToLayers:
             for layer in ops2layers.layers:
                 qco_by_opset_name = _current_tpc.get().tp_model.get_config_options_by_operators_set(ops2layers.name)
                 if layer in existing_layers:
-                    raise Exception(f'Found layer {layer.__name__} in more than one '
-                                    f'OperatorsSet')
+                    Logger.error(f'Found layer {layer.__name__} in more than one '
+                                 f'OperatorsSet')  # pragma: no cover
                 else:
                     existing_layers.update({layer: qco_by_opset_name})
-
-

model_compression_toolkit/core/common/target_platform/targetplatform2framework/target_platform_capabilities.py

@@ -131,7 +131,8 @@ class TargetPlatformCapabilities(ImmutableClass):
         if isinstance(tpc_component, OperationsSetToLayers):
             self.op_sets_to_layers += tpc_component
         else:
-            raise Exception(f'Trying to append an unfamiliar TargetPlatformCapabilitiesComponent of type: {type(tpc_component)}')
+            Logger.error(f'Trying to append an unfamiliar TargetPlatformCapabilitiesComponent of type: '
+                         f'{type(tpc_component)}')  # pragma: no cover
 
     def __enter__(self):
         """
@@ -175,7 +176,7 @@ class TargetPlatformCapabilities(ImmutableClass):
             QuantizationConfigOptions of the node.
         """
         if node is None:
-            raise Exception(f'Can not retrieve QC options for None node')
+            Logger.error(f'Can not retrieve QC options for None node')  # pragma: no cover
         for fl, qco in self.filterlayer2qco.items():
             if fl.match(node):
                 return qco
@@ -205,7 +206,6 @@ class TargetPlatformCapabilities(ImmutableClass):
                     layer2qco.update({l: qco})
         return layer2qco, filterlayer2qco
 
-
     def remove_fusing_names_from_not_used_list(self):
         """
         Remove OperatorSets names from the list of the unused sets (so a warning
@@ -235,5 +235,3 @@ class TargetPlatformCapabilities(ImmutableClass):
         """
         for op in self.__tp_model_opsets_not_used:
             Logger.warning(f'{op} is defined in TargetPlatformModel, but is not used in TargetPlatformCapabilities.')
-
-

model_compression_toolkit/core/keras/back2framework/instance_builder.py

@@ -20,38 +20,25 @@ from typing import List, Dict, Callable
 from networkx.algorithms.dag import topological_sort
 
 import tensorflow as tf
-from tensorflow.keras.layers import Layer
+from tensorflow.keras.layers import Layer, InputLayer
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.keras.constants import LAYER_NAME
 
 
-def identity_wrapper(node: BaseNode, layer: Layer):
-    """
-    A function which takes a computational graph node and a keras layer and return an identity wrapping which return the layer itself
-    Args:
-        node: A node of mct graph.
-        layer: A keras layer
-
-    Returns: keras layer
-
-    """
-    return layer
-
-
 class OperationHandler:
     """
     Class to handle conversions from graph nodes to Keras operators and retrieving them.
     """
 
-    def __init__(self, graph: Graph, wrapper: Callable = identity_wrapper):
+    def __init__(self, graph: Graph):
         # hold nodes after sorting them
         self.node_sort = list(topological_sort(graph))
 
         self.layer_to_node_dict = {}
 
         # hold dictionary from node to its equivalent Keras layer
-        self.node_to_fw_op_dict = instance_builder(self.node_sort, wrapper)
+        self.node_to_fw_op_dict = instance_builder(self.node_sort)
 
     def get_node_op_function(self, n: BaseNode) -> Layer:
         """
@@ -86,10 +73,15 @@ def node_builder(n: common.BaseNode) -> Layer:
     Returns:
         Keras layer that was built from the node.
     """
-
     framework_attr = copy.copy(n.framework_attr)
+    if n.layer_class is InputLayer:
+        # replace input node with identity, so can wrap it with QuantizationWrapper
+        _layer_class = Layer  # Identity
+        framework_attr = {}
+    else:
+        _layer_class = n.layer_class
     framework_attr[LAYER_NAME] = n.name  # Overwrite framework name to identical graph node name
-    node_instance = n.layer_class.from_config(framework_attr)  # Build layer from node's configuration.
+    node_instance = _layer_class.from_config(framework_attr)  # Build layer from node's configuration.
     with tf.name_scope(n.name):
         # Add layer name to default weight name to avoid name duplications
         node_instance.build(n.input_shape)
@@ -98,13 +90,12 @@ def node_builder(n: common.BaseNode) -> Layer:
     return node_instance
 
 
-def instance_builder(toposort: List[BaseNode], wrapper: Callable) -> Dict[BaseNode, Layer]:
+def instance_builder(toposort: List[BaseNode]) -> Dict[BaseNode, Layer]:
     """
     Build a dictionary of nodes to their corresponding Keras
     layers, given a list of nodes.
 
     Args:
-        wrapper: A function wrapper keras Layers.
         toposort: List of nodes sorted topological to build their layers.
 
     Returns:
@@ -114,7 +105,7 @@ def instance_builder(toposort: List[BaseNode], wrapper: Callable) -> Dict[BaseNo
     nodes_dict = dict()
     for n in toposort:
         if not n.reuse:  # Hold a single node in dictionary for all reused nodes from the same layer.
-            keras_node = wrapper(n, node_builder(n))
+            keras_node = node_builder(n)
             nodes_dict.update({n: keras_node})
 
     return nodes_dict