mct-nightly 1.1.0.7012022.post2611__py3-none-any.whl → 1.1.0.07122021-002414__py3-none-any.whl

model_compression_toolkit/keras/graph_substitutions/substitutions/remove_relu_upper_bound.py

@@ -17,7 +17,7 @@
 from tensorflow.keras.layers import ReLU
 
 from model_compression_toolkit import common
-from model_compression_toolkit.common import Graph, BaseNode
+from model_compression_toolkit.common import Graph, Node
 from model_compression_toolkit.common.graph.graph_matchers import NodeOperationMatcher,NodeFrameworkAttrMatcher
 from model_compression_toolkit.keras.constants import RELU_MAX_VALUE
 from model_compression_toolkit.common.constants import THRESHOLD
@@ -40,7 +40,7 @@ class RemoveReLUUpperBound(common.BaseSubstitution):
 
     def substitute(self,
                    graph: Graph,
-                   node: BaseNode) -> Graph:
+                   node: Node) -> Graph:
         """
         Remove ReLU upper bound if its activation threshold bounds it anyway at
         the same value.

model_compression_toolkit/keras/graph_substitutions/substitutions/scale_equalization.py

@@ -21,7 +21,7 @@ import numpy as np
 from tensorflow.keras.layers import DepthwiseConv2D, Conv2D, Dense, Conv2DTranspose, Activation, ReLU, ZeroPadding2D
 
 from model_compression_toolkit import common
-from model_compression_toolkit.common import Graph, BaseNode
+from model_compression_toolkit.common import Graph, Node
 from model_compression_toolkit.common.constants import OUTPUT_SCALE, THRESHOLD
 from model_compression_toolkit.common.defaultdict import DefaultDict
 from model_compression_toolkit.common.framework_info import FrameworkInfo
@@ -59,7 +59,7 @@ MATCHER_MID_WITH_PAD = WalkMatcher([op2d_node, mid_activation_nodes, zeropad_nod
 
 
 def scale_reshaping(scale: np.ndarray,
-                    op2d: common.BaseNode,
+                    op2d: common.Node,
                     kernel_channel_mapping: DefaultDict,
                     in_channels: bool = True) -> np.ndarray:
     """
@@ -87,8 +87,8 @@ def scale_reshaping(scale: np.ndarray,
 def update_linear_nodes(graph:Graph,
                         qc: QuantizationConfig,
                         fw_info: FrameworkInfo,
-                        first_op2d_node: BaseNode,
-                        second_op2d_node: BaseNode,
+                        first_op2d_node: Node,
+                        second_op2d_node: Node,
                         scale_factor: np.ndarray):
     """
     Scale the weights of two linear nodes with a scale factor. Each node is scaled in
@@ -132,7 +132,7 @@ def update_linear_nodes(graph:Graph,
 
 
 def calculate_scale_correction(graph: Graph,
-                               activation_node: BaseNode,
+                               activation_node: Node,
                                eps: float = 1e-6) -> tuple:
     """
     Compute a scale factor by the activation node threshold and its outputs statistics in
@@ -172,9 +172,9 @@ def calculate_scale_correction(graph: Graph,
 def scale_equalization_lnl(graph: Graph,
                            qc: QuantizationConfig,
                            fw_info: FrameworkInfo,
-                           first_op2d_node: BaseNode,
-                           n_node: BaseNode,
-                           second_op2d_node: BaseNode):
+                           first_op2d_node: Node,
+                           n_node: Node,
+                           second_op2d_node: Node):
     """
     Compute a scale factor to scale all activation node's outputs such that
     its maximum per-channel is the constrained threshold of the activation node.
@@ -235,7 +235,7 @@ class BaseScaleEqualization(common.BaseSubstitution):
 
     def substitute(self,
                    graph: Graph,
-                   nodes_list: List[BaseNode]) -> Graph:
+                   nodes_list: List[Node]) -> Graph:
         """
         Scale each channel of the weights of two linear nodes,
         in order to use the entire constrained range when activations are quantized.

model_compression_toolkit/keras/graph_substitutions/substitutions/separableconv_decomposition.py

@@ -19,7 +19,7 @@ from tensorflow.keras.layers import SeparableConv2D, Conv2D, DepthwiseConv2D
 from model_compression_toolkit import common
 from model_compression_toolkit.common.graph.base_graph import Graph
 from model_compression_toolkit.common.graph.graph_matchers import NodeOperationMatcher
-from model_compression_toolkit.common.graph.base_node import BaseNode
+from model_compression_toolkit.common.graph.node import Node
 from model_compression_toolkit.keras.constants import KERNEL, DEPTHWISE_KERNEL, BIAS, KERNEL_SIZE, PADDING, \
     STRIDES, USE_BIAS, LINEAR, ACTIVATION, TRAINABLE, FILTERS, PAD_VALID
 
@@ -57,7 +57,7 @@ class SeparableConvDecomposition(common.BaseSubstitution):
 
     def substitute(self,
                    graph: Graph,
-                   separable_node: BaseNode) -> Graph:
+                   separable_node: Node) -> Graph:
         """
         Remove a SeparableConv2D node from the graph, and replace it with two equivalent nodes: DepthwiseConv2D
         and Conv2D. The SeparableConv2D attributes are split to relevant attributes for each node.
@@ -114,28 +114,28 @@ class SeparableConvDecomposition(common.BaseSubstitution):
         dw_node_name = separable_node.name + '_dw' if not separable_node.reuse else '_'.join(separable_node.name.split('_')[:-2]) + '_dw_' + '_'.join(separable_node.name.split('_')[-2:])
 
         # create new nodes
-        dw_node = common.graph.BaseNode(dw_node_name,
-                                        dw_framework_attr,
-                                        separable_node.input_shape,
-                                        dw_output_shape,
-                                        dw_weights_dict,
-                                        dw_layer_class,
-                                        reuse=separable_node.reuse,
-                                        reuse_group=separable_node.reuse_group)
+        dw_node = common.graph.Node(dw_node_name,
+                                    dw_framework_attr,
+                                    separable_node.input_shape,
+                                    dw_output_shape,
+                                    dw_weights_dict,
+                                    dw_layer_class,
+                                    reuse=separable_node.reuse,
+                                    reuse_group=separable_node.reuse_group)
 
         # If the SeparableConv2D is reused, we need to keep the pointwise node as reused as well,
         # so we keep the names convention with adding the suffix of "_reuse_X".
         pw_node_name = separable_node.name + '_pw' if not separable_node.reuse else '_'.join(separable_node.name.split('_')[:-2]) + '_pw_' + '_'.join(separable_node.name.split('_')[-2:])
 
-        pw_node = common.graph.BaseNode(pw_node_name,
-                                        pw_framework_attr,
-                                        pw_input_shape,
-                                        separable_node.output_shape,
-                                        pw_weights_dict,
-                                        pw_layer_class,
-                                        reuse=separable_node.reuse,
-                                        reuse_group=separable_node.reuse_group
-                                        )
+        pw_node = common.graph.Node(pw_node_name,
+                                    pw_framework_attr,
+                                    pw_input_shape,
+                                    separable_node.output_shape,
+                                    pw_weights_dict,
+                                    pw_layer_class,
+                                    reuse=separable_node.reuse,
+                                    reuse_group=separable_node.reuse_group
+                                    )
 
         graph.add_node(dw_node)
         graph.add_node(pw_node)

model_compression_toolkit/keras/graph_substitutions/substitutions/shift_negative_activation.py

@@ -16,8 +16,6 @@
 import tensorflow as tf
 
 # As from Tensorflow 2.6, keras is a separate package and some classes should be imported differently.
-from model_compression_toolkit.common.framework_implementation import FrameworkImplementation
-
 if tf.__version__ < "2.6":
     from tensorflow.python.keras.engine.base_layer import TensorFlowOpLayer
 else:
@@ -31,14 +29,13 @@ from tensorflow.keras.layers import Activation, Conv2D, Dense, DepthwiseConv2D,
 from typing import Tuple, Any
 
 from model_compression_toolkit import common
-from model_compression_toolkit.common import FrameworkInfo, Graph, BaseNode
-from model_compression_toolkit.common.constants import FLOAT_32, DATA_TYPE, THRESHOLD, SIGNED
+from model_compression_toolkit.common import FrameworkInfo, Graph, Node
+from model_compression_toolkit.common.constants import FLOAT_32, DATA_TYPE, THRESHOLD
 from model_compression_toolkit.common.graph.graph_matchers import EdgeMatcher
 from model_compression_toolkit.common.graph.graph_matchers import NodeOperationMatcher, \
     NodeFrameworkAttrMatcher
 
-from model_compression_toolkit.common.quantization.set_node_quantization_config import create_node_activation_qc, \
-    set_quantization_configs_to_node
+from model_compression_toolkit.common.quantization.set_node_quantization_config import create_node_activation_qc
 from model_compression_toolkit.common.quantization.quantization_config import QuantizationConfig
 from model_compression_toolkit.common.quantization.quantization_params_generation.qparams_activations_computation \
     import \
@@ -103,7 +100,7 @@ PAD_NODE = NodeOperationMatcher(ZeroPadding2D)
 
 def create_add_node(add_value: float,
                     prev_node_name: str,
-                    input_shape: tuple) -> BaseNode:
+                    input_shape: tuple) -> Node:
     """
     Create a new Add node, with a constant to add.
     The name of the node is determined by its previous node's name.
@@ -133,13 +130,13 @@ def create_add_node(add_value: float,
         CONSTANTS: {1: np.array([[[[add_value]]]],
                                 dtype=np.float32)}}
 
-    add_node = common.graph.BaseNode(add_node_name,
-                                     fw_attr,
-                                     input_shape,
-                                     input_shape,
-                                     weights={},
-                                     quantization_attr={},
-                                     layer_class=TensorFlowOpLayer)
+    add_node = common.graph.Node(add_node_name,
+                                 fw_attr,
+                                 input_shape,
+                                 input_shape,
+                                 weights={},
+                                 quantization_attr={},
+                                 layer_class=TensorFlowOpLayer)
     return add_node
 
 
@@ -150,7 +147,7 @@ def create_pad_node(next_node_name: str,
                     pad_top: int,
                     pad_btm: int,
                     pad_left: int,
-                    pad_right: int) -> BaseNode:
+                    pad_right: int) -> Node:
     """
     Create a pad node with a constant value to pad its input tensor.
 
@@ -192,17 +189,17 @@ def create_pad_node(next_node_name: str,
     padded_shape = list(input_shape)
     padded_shape[1] += pad_top + pad_btm
     padded_shape[2] += pad_left + pad_right
-    pad_node = common.graph.BaseNode(pad_node_name,
-                                     fw_attr,
-                                     input_shape,
-                                     tuple(padded_shape),
-                                     weights={},
-                                     quantization_attr={},
-                                     layer_class=TensorFlowOpLayer)
+    pad_node = common.graph.Node(pad_node_name,
+                                 fw_attr,
+                                 input_shape,
+                                 tuple(padded_shape),
+                                 weights={},
+                                 quantization_attr={},
+                                 layer_class=TensorFlowOpLayer)
     return pad_node
 
 
-def compute_op2d_padding(op2d_node: BaseNode) -> Tuple[int, int, int, int]:
+def compute_op2d_padding(op2d_node: Node) -> Tuple[int, int, int, int]:
     """
     Compute the padding around an input tensor of a linear node.
     This is needed to replace tensorflow 'same' padding with actual number of elements to pad.
@@ -231,7 +228,7 @@ def compute_op2d_padding(op2d_node: BaseNode) -> Tuple[int, int, int, int]:
     return pad_top, pad_btm, pad_left, pad_right
 
 
-def op2d_bias_correction(op2d_node: common.BaseNode,
+def op2d_bias_correction(op2d_node: common.Node,
                          shift_to_correct: float):
     """
     Compute the correction term to add to the op2d node's bias
@@ -269,9 +266,9 @@ def op2d_bias_correction(op2d_node: common.BaseNode,
 
 
 def insert_node_between_two_nodes(graph: Graph,
-                                  node_to_insert: BaseNode,
-                                  first_node: BaseNode,
-                                  last_node: BaseNode):
+                                  node_to_insert: Node,
+                                  first_node: Node,
+                                  last_node: Node):
     """
     Insert a new node in a graph between two nodes.
 
@@ -293,8 +290,8 @@ def insert_node_between_two_nodes(graph: Graph,
 
 
 def insert_node_after_node(graph: Graph,
-                           node_to_insert: BaseNode,
-                           first_node: BaseNode):
+                           node_to_insert: Node,
+                           first_node: Node):
     """
     Insert a new node to a graph after an existing node in the graph.
     Check before insertion that the node (that we add the new node after) has
@@ -316,8 +313,8 @@ def insert_node_after_node(graph: Graph,
 
 
 def insert_node_before_node(graph: Graph,
-                            node_to_insert: BaseNode,
-                            last_node: BaseNode):
+                            node_to_insert: Node,
+                            last_node: Node):
     """
     Insert a new node to a graph before an existing node in the graph.
     Check before insertion that the node (that we add the new node before) has
@@ -338,9 +335,9 @@ def insert_node_before_node(graph: Graph,
 
 
 def remove_node_between_two_nodes(graph: Graph,
-                                  node_to_remove: BaseNode,
-                                  first_node: BaseNode,
-                                  last_node: BaseNode):
+                                  node_to_remove: Node,
+                                  first_node: Node,
+                                  last_node: Node):
     """
     Remove a node from a graph and connect its previous node to
     its next node after the removal.
@@ -363,12 +360,12 @@ def remove_node_between_two_nodes(graph: Graph,
     graph.remove_node(node_to_remove)
 
 
-def shift_negative_function(graph: Graph,
-                            qc: QuantizationConfig,
-                            non_linear_node: BaseNode,
-                            op2d_node: BaseNode,
+def shift_negative_function(graph,
+                            qc,
+                            non_linear_node,
+                            op2d_node,
                             fw_info: FrameworkInfo,
-                            zero_padding_node: BaseNode = None) -> Graph:
+                            zero_padding_node=None):
     """
     Shift the output of a non-linear activation by its minimal output value (quantized) such
     that all values after the shifting are positive.
@@ -457,15 +454,6 @@ def shift_negative_function(graph: Graph,
                                    add_node.output_shape,
                                    pad_top, pad_btm, pad_left, pad_right)
 
-        # Set quantization configuration to node, even though we do not quantize it:
-        set_quantization_configs_to_node(fw_info=fw_info,
-                                         node=pad_node,
-                                         quant_config=qc)
-
-        pad_node.activation_quantization_cfg.enable_activation_quantization = False
-        for weight_qc in pad_node.candidates_weights_quantization_cfg:
-            weight_qc.enable_weights_quantization = False
-
         # Insert a pad node between the add node to the op2d, and create statistics for the pad node
         insert_node_before_node(graph,
                                 node_to_insert=pad_node,
@@ -476,32 +464,25 @@ def shift_negative_function(graph: Graph,
 
         op2d_node.input_shape = pad_node.output_shape
 
-    set_quantization_configs_to_node(fw_info=fw_info,
-                                     node=add_node,
-                                     quant_config=qc)
-
-    add_node.activation_quantization_cfg.enable_activation_quantization = False
-
-    for weight_qc in add_node.candidates_weights_quantization_cfg:
-        weight_qc.enable_weights_quantization = False
-
     add_node.activation_quantization_cfg = create_node_activation_qc(qc,
-                                                                     fw_info)
+                                                                     fw_info,
+                                                                     add_node_stats_collector.use_min_max)
 
-    add_node.activation_quantization_cfg.set_activation_quantization_param({THRESHOLD: activation_threshold,
-                                                                            SIGNED: False})
+    add_node.activation_quantization_cfg.set_activation_quantization_param({THRESHOLD: activation_threshold})
+    add_node.activation_quantization_cfg.activation_is_signed = False
 
     if non_linear_node.activation_quantization_cfg.shift_negative_threshold_recalculation:
-        activation_param = get_activations_qparams(add_node, graph)
-        assert activation_param.get(SIGNED) == False
+        activation_param, activation_is_signed = get_activations_qparams(add_node, graph)
+        assert activation_is_signed == False
         add_node.activation_quantization_cfg.set_activation_quantization_param(activation_param)
+        add_node.activation_quantization_cfg.activation_is_signed = False
 
     return graph
 
 
-def get_next_nodes_to_correct(n: BaseNode,
+def get_next_nodes_to_correct(n: Node,
                               graph: Graph,
-                              pad_node_to_consider: BaseNode = None) -> Tuple[Any, Any]:
+                              pad_node_to_consider: Node = None) -> Tuple[Any, Any]:
     """
     Search for the next linear node of a given node. Go over
     the next nodes of the node and recursively search for a linear node.

model_compression_toolkit/keras/keras_implementation.py

@@ -6,11 +6,9 @@ from tensorflow.keras.models import Model
 
 from model_compression_toolkit import QuantizationConfig, FrameworkInfo, common, GradientPTQConfig, \
     MixedPrecisionQuantizationConfig
-from model_compression_toolkit.common import Graph, BaseNode
-from model_compression_toolkit.common.collectors.statistics_collector import BaseStatsCollector
+from model_compression_toolkit.common import Graph, Node
 from model_compression_toolkit.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.common.model_builder_mode import ModelBuilderMode
-from model_compression_toolkit.common.node_prior_info import NodePriorInfo
 from model_compression_toolkit.common.user_info import UserInformation
 from model_compression_toolkit.keras.back2framework.model_builder import model_builder
 from model_compression_toolkit.keras.default_framework_info import DEFAULT_KERAS_INFO
@@ -32,10 +30,9 @@ from model_compression_toolkit.keras.graph_substitutions.substitutions.separable
     SeparableConvDecomposition
 from model_compression_toolkit.keras.graph_substitutions.substitutions.shift_negative_activation import \
     apply_shift_negative_correction
-from model_compression_toolkit.keras.keras_node_prior_info import create_node_prior_info
 from model_compression_toolkit.keras.mixed_precision.sensitivity_evaluation import get_sensitivity_evaluation
 from model_compression_toolkit.keras.reader.reader import model_reader
-from model_compression_toolkit.common.collectors.statistics_collector_generator import create_stats_collector_for_node
+from model_compression_toolkit.keras.tensor_marking import get_node_stats_collector
 import model_compression_toolkit.keras.constants as keras_constants
 
 
@@ -121,22 +118,21 @@ class KerasImplementation(FrameworkImplementation):
                                                qc,
                                                fw_info)
 
-    def attach_sc_to_node(self,
-                          node: BaseNode,
-                          output_channel_index: int) -> BaseStatsCollector:
+    def attach_sc_to_node(self, node: Node,
+                          fw_info: FrameworkInfo) -> common.statistics_collector.BaseStatsContainer:
         """
         Return a statistics collector that should be attached to a node's output
         during statistics collection.
 
         Args:
             node: Node to return its collector.
-            output_channel_index: Index of output channels of layers in the model's framework.
+            fw_info: FrameworkInfo object with information about the specific framework's model
 
         Returns:
             Statistics collector for the node.
         """
-        return create_stats_collector_for_node(node,
-                                               output_channel_index=output_channel_index)
+        return get_node_stats_collector(node,
+                                        fw_info)
 
     def get_substitutions_marking(self) -> List[common.BaseSubstitution]:
         """
@@ -231,6 +227,7 @@ class KerasImplementation(FrameworkImplementation):
                                      gptq_config,
                                      fw_info)
 
+
     def get_sensitivity_evaluation_fn(self,
                                       graph: Graph,
                                       quant_config: MixedPrecisionQuantizationConfig,
@@ -257,20 +254,3 @@ class KerasImplementation(FrameworkImplementation):
                                           metrics_weights,
                                           representative_data_gen,
                                           fw_info)
-
-    def get_node_prior_info(self,
-                            node: BaseNode,
-                            fw_info: FrameworkInfo) -> NodePriorInfo:
-        """
-        Get a NodePriorInfo object for a node that represents a Keras layer.
-
-        Args:
-            node: Node to get its prior info.
-            fw_info: Framework specific information needed to create the prior info of the node.
-
-        Returns:
-            NodePriorInfo with information about the node.
-        """
-
-        return create_node_prior_info(node=node,
-                                      fw_info=fw_info)

model_compression_toolkit/keras/mixed_precision/sensitivity_evaluation.py

@@ -20,7 +20,7 @@ from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapp
 from typing import Callable, List, Any
 
 from model_compression_toolkit.common.framework_info import FrameworkInfo
-from model_compression_toolkit.common import BaseNode
+from model_compression_toolkit.common import Node
 from model_compression_toolkit.common.graph.base_graph import Graph
 from model_compression_toolkit.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfig
@@ -207,7 +207,7 @@ def _build_distance_matrix(baseline_tensors: List[Tensor],
 
 
 def _build_baseline_model(graph: Graph,
-                          interest_points: List[BaseNode]) -> Model:
+                          interest_points: List[Node]) -> Model:
     """
     Build a Keras baseline model to compare inferences of the MP model to.
     The baseline model is the float model we build from the graph.

model_compression_toolkit/keras/quantization_facade.py

@@ -33,7 +33,6 @@ if importlib.util.find_spec("tensorflow") is not None\
         and importlib.util.find_spec("tensorflow_model_optimization") is not None:
     from model_compression_toolkit.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.keras.keras_implementation import KerasImplementation
-    from model_compression_toolkit.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
 
     def keras_post_training_quantization(in_model: Model,
@@ -85,8 +84,7 @@ if importlib.util.find_spec("tensorflow") is not None\
             >>> quantized_model, quantization_info = mct.keras_post_training_quantization(model, repr_datagen)
 
         """
-        KerasModelValidation(model=in_model,
-                             fw_info=fw_info).validate()
+
         return post_training_quantization(in_model,
                                           representative_data_gen,
                                           n_iter,
@@ -168,8 +166,6 @@ if importlib.util.find_spec("tensorflow") is not None\
              For more configuration options, please take a look at our `API documentation <https://sony.github.io/model_optimization/api/api_docs/modules/mixed_precision_quantization_config.html>`_.
 
          """
-        KerasModelValidation(model=in_model,
-                             fw_info=fw_info).validate()
 
         if target_kpi is None:
             common.Logger.warning("No KPI was passed. Using non mixed-precision compression process...")

model_compression_toolkit/keras/quantizer/fake_quant_builder.py

@@ -21,7 +21,7 @@ import tensorflow as tf
 import numpy as np
 from tensorflow.python.util.object_identity import Reference as TFReference
 
-from model_compression_toolkit.common.constants import THRESHOLD, SIGNED
+from model_compression_toolkit.common.constants import THRESHOLD
 
 
 def quantizer_min_max_calculator(threshold: np.ndarray,
@@ -52,6 +52,7 @@ def quantizer_min_max_calculator(threshold: np.ndarray,
 
 
 def constraint_quantization(activation_n_bits: int,
+                            activation_is_signed: bool,
                             quantization_params: dict) -> Callable:
     """
     Use a NodeQuantizationConfig to compute a quantizer min/max values, and use it to
@@ -59,15 +60,14 @@ def constraint_quantization(activation_n_bits: int,
 
     Args:
         activation_n_bits: Number of bits to use for quantization.
+        activation_is_signed: Whether the quantization range should include negative values or not.
         quantization_params: Dictionary of specific parameters for this quantization function.
 
     Returns:
         A fake quantization node.
     """
     activation_threshold = quantization_params.get(THRESHOLD)
-    activation_is_signed = quantization_params.get(SIGNED)
-
-    if activation_threshold is None or activation_is_signed is None:
+    if activation_threshold is None:
         return None
 
     min_value, max_value = quantizer_min_max_calculator(activation_threshold,

model_compression_toolkit/keras/quantizer/gradient_ptq/activation_quantizer.py

@@ -23,7 +23,7 @@ from model_compression_toolkit.keras.quantizer.base_quantizer import BaseTrainab
 from model_compression_toolkit.keras.quantizer.gradient_ptq.utils import symmetric_quantizer
 from model_compression_toolkit.keras.quantizer.gradient_ptq.utils import ste_round
 from model_compression_toolkit import common
-from model_compression_toolkit.common.constants import THRESHOLD, SIGNED
+from model_compression_toolkit.common.constants import THRESHOLD
 
 
 class TrainableQuantizer(BaseTrainableQuantizer):
@@ -171,8 +171,7 @@ class TrainableQuantizer(BaseTrainableQuantizer):
 
         threshold_change = np.asarray(new_threshold / old_threshold).flatten()
         common.Logger.info(f"Layer '{layer.layer.name}' has total threshold change of {str(threshold_change)}")
-        return {THRESHOLD: new_threshold.numpy().reshape(self.threshold_shape),
-                SIGNED: self.signed}
+        return {THRESHOLD: new_threshold.numpy().reshape(self.threshold_shape)}
 
     def get_trainable_parameters(self):
         """

model_compression_toolkit/keras/quantizer/gradient_ptq/activation_quantizer_gptq_config.py

@@ -17,15 +17,10 @@ from typing import List, Any, Dict
 
 from tensorflow.python.training.tracking.data_structures import ListWrapper
 
-from model_compression_toolkit.common.constants import THRESHOLD, SIGNED
+from model_compression_toolkit.common.constants import THRESHOLD
 from model_compression_toolkit.keras.quantizer.gradient_ptq.activation_quantizer import TrainableQuantizer
 from model_compression_toolkit.keras.quantizer.gradient_ptq.base_quantizer_gptq_config import BaseQuantizeConfig
-import tensorflow as tf
-# As from Tensorflow 2.6, keras is a separate package and some classes should be imported differently.
-if tf.__version__ < "2.6":
-    from tensorflow.python.keras.layers import Layer
-else:
-    from keras.engine.base_layer import Layer
+from tensorflow.python.keras.layers import Layer
 from tensorflow import Tensor
 from tensorflow_model_optimization.python.core.quantization.keras.quantizers import Quantizer
 
@@ -37,16 +32,17 @@ class ActivationQuantizeConfig(BaseQuantizeConfig):
 
     def __init__(self,
                  activation_quantization_params: dict,
+                 signed: bool,
                  num_bits: int = 8):
         """
         Initialize a TrainableQuantizer and set as the activation quantizer.
 
         Args:
             activation_quantization_params: Parameters to use for quantization.
+            signed: Quantization range is signed or unsigned.
             num_bits: Number of bits to use for quantization.
         """
         threshold_values = activation_quantization_params.get(THRESHOLD)
-        signed = activation_quantization_params.get(SIGNED)
         self.activation_quantizer = TrainableQuantizer(num_bits=num_bits,
                                                        per_axis=False,
                                                        threshold_values=threshold_values,

model_compression_toolkit/keras/quantizer/gradient_ptq/activation_weight_quantizer_gptq_config.py

@@ -16,16 +16,11 @@
 from typing import List, Tuple, Any, Dict
 
 from tensorflow import Tensor
-import tensorflow as tf
-# As from Tensorflow 2.6, keras is a separate package and some classes should be imported differently.
-if tf.__version__ < "2.6":
-    from tensorflow.python.keras.layers import Layer
-else:
-    from keras.engine.base_layer import Layer
+from tensorflow.python.keras.layers import Layer
 from tensorflow.python.training.tracking.data_structures import ListWrapper
 from tensorflow_model_optimization.python.core.quantization.keras.quantizers import Quantizer
 
-from model_compression_toolkit.common.constants import THRESHOLD, SIGNED
+from model_compression_toolkit.common.constants import THRESHOLD
 from model_compression_toolkit.keras.quantizer.gradient_ptq.weight_quantizer import TrainableWeightQuantizer
 from model_compression_toolkit.keras.quantizer.gradient_ptq.activation_quantizer import TrainableQuantizer
 from model_compression_toolkit.keras.quantizer.gradient_ptq.base_quantizer_gptq_config import BaseQuantizeConfig
@@ -44,6 +39,7 @@ class ActivationAndWeightQuantizeConfig(BaseQuantizeConfig):
                  weight_channel_axis: int,
                  weight_num_bits: int,
                  activation_quantization_params: dict,
+                 activation_signed: bool,
                  activation_num_bits: int = 8,
                  max_lsbs_change: int = 8):
         """
@@ -54,13 +50,12 @@ class ActivationAndWeightQuantizeConfig(BaseQuantizeConfig):
             weight_channel_axis: Channel index to quantize when quantizing the weight per-channel.
             weight_num_bits: Number of bits to use for weight quantization.
             activation_quantization_params: Parameters to use for the activation quantization.
+            activation_signed: Quantization range is signed or unsigned.
             activation_num_bits: Number of bits to use for quantization of the activation.
 
         """
 
         activation_threshold_values = activation_quantization_params.get(THRESHOLD)
-        activation_signed = activation_quantization_params.get(SIGNED)
-
         self.activation_quantizer = TrainableQuantizer(num_bits=activation_num_bits,
                                                        per_axis=False,
                                                        threshold_values=activation_threshold_values,