mct-nightly 1.1.0.6012022.post2521__py3-none-any.whl → 1.1.0.07122021-002414__py3-none-any.whl

model_compression_toolkit/common/{collectors/statistics_collector.py → statistics_collector.py}

@@ -20,17 +20,20 @@ from typing import Any, Tuple
 
 import numpy as np
 
-from model_compression_toolkit.common.framework_info import FrameworkInfo, ChannelAxis
 from model_compression_toolkit.common.collectors.histogram_collector import HistogramCollector
 from model_compression_toolkit.common.collectors.mean_collector import MeanCollector
 from model_compression_toolkit.common.collectors.min_max_per_channel_collector import MinMaxPerChannelCollector
 
 
-class BaseStatsCollector(object):
+class BaseStatsContainer(object):
     """
-    Base class for statistics collection (contains multiple collectors such as mean collector,
+    Base class for statistics collection container (contain multiple statistics collector such as mean collector,
     histogram collector, etc.).
     """
+    def __init__(self):
+        # Disable histogram collection. Enable in specific collectors if needed
+        self.collect_histogram = False
+        self.use_min_max = False
 
     def require_collection(self) -> bool:
         """
@@ -50,13 +53,12 @@ class BaseStatsCollector(object):
         raise Exception(f'update_statistics is not implemented in {self.__class__.__name__}')
 
 
-class StatsCollector(BaseStatsCollector):
+class StatsContainer(BaseStatsContainer):
     """
     Class to wrap all statistics that are being collected for an input/output node.
     """
 
     def __init__(self,
-                 output_channel_index: ChannelAxis,
                  init_min_value: float = None,
                  init_max_value: float = None):
         """
@@ -64,17 +66,18 @@ class StatsCollector(BaseStatsCollector):
         Set initial min/max values if are known.
 
         Args:
-            output_channel_index: Index of output channels.
             init_min_value: Initial min value for min/max stored values.
             init_max_value: Initial max value for min/max stored values.
         """
 
         super().__init__()
-        self.hc = HistogramCollector()
-        self.mc = MeanCollector(axis=output_channel_index)
+        self.use_min_max = is_number(init_min_value) and is_number(init_max_value)
+        self.collect_histogram = True
+        if self.collect_histogram:
+            self.hc = HistogramCollector()
+        self.mc = MeanCollector()
         self.mpcc = MinMaxPerChannelCollector(init_min_value=init_min_value,
-                                              init_max_value=init_max_value,
-                                              axis=output_channel_index)
+                                              init_max_value=init_max_value)
 
     def update_statistics(self, x: Any):
         """
@@ -85,7 +88,8 @@ class StatsCollector(BaseStatsCollector):
         """
 
         x = standardize_tensor(x)
-        self.hc.update(x)
+        if self.collect_histogram:
+            self.hc.update(x)
         self.mc.update(x)
         self.mpcc.update(x)
 
@@ -139,7 +143,7 @@ class StatsCollector(BaseStatsCollector):
         return True
 
 
-class NoStatsCollector(BaseStatsCollector):
+class NoStatsContainer(BaseStatsContainer):
     """
     Class that inherits from base tensor.
     Indicating that for a point in a graph we should not gather statistics.
@@ -207,51 +211,51 @@ def standardize_tensor(x: Any) -> np.ndarray:
     return x
 
 
-def shift_statistics(collector: BaseStatsCollector,
-                     shift_value: np.ndarray) -> BaseStatsCollector:
+def shift_statistics(collector: BaseStatsContainer,
+                     shift_value: np.ndarray) -> BaseStatsContainer:
     """
-    Shift all statistics in collectors of a statistics collector by a
+    Shift all statistics in collectors of a statistics container by a
     value (or a value per-channel).
 
     Args:
-        collector: Statistics collector to shift its collectors.
+        collector: Statistics container to shift its collectors.
         shift_value: Value to shift all statistics by.
 
     Returns:
-        New copy of the collector with shifted statistics.
+        New copy of the container with shifted statistics.
 
     """
 
     shifted_collector = deepcopy(collector)
-    if isinstance(collector, StatsCollector):
+    if isinstance(collector, StatsContainer):
         shifted_collector.mpcc.shift(shift_value)
         shifted_collector.mc.shift(shift_value)
-        if shifted_collector.require_collection():
+        if shifted_collector.collect_histogram:
             shifted_collector.hc.shift(shift_value)
 
     return shifted_collector
 
 
-def scale_statistics(collector: BaseStatsCollector,
-                     scale_value: np.ndarray) -> BaseStatsCollector:
+def scale_statistics(collector: BaseStatsContainer,
+                     scale_value: np.ndarray) -> BaseStatsContainer:
     """
-    Scale all statistics in collectors of a statistics collector
+    Scale all statistics in collectors of a statistics container
     by a factor (or a factor per-channel).
 
     Args:
-        collector: Statistics collector to shift its collectors.
+        collector: Statistics container to shift its collectors.
         scale_value: Value to shift all statistics by.
 
     Returns:
-        New copy of the collector with scaled statistics.
+        New copy of the container with scaled statistics.
 
     """
 
     scaled_collector = deepcopy(collector)
-    if isinstance(collector, StatsCollector):
+    if isinstance(collector, StatsContainer):
         scaled_collector.mpcc.scale(scale_value)
         scaled_collector.mc.scale(scale_value)
-        if scaled_collector.require_collection():
+        if scaled_collector.collect_histogram:
             scaled_collector.hc.scale(scale_value)
 
     return scaled_collector

model_compression_toolkit/common/visualization/tensorboard_writer.py

@@ -29,10 +29,10 @@ from tensorboard.compat.proto.summary_pb2 import HistogramProto
 from tensorboard.compat.proto.summary_pb2 import Summary
 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 typing import List, Any, Dict, Callable
 
-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.statistics_collector import BaseStatsContainer
 
 DEVICE_STEP_STATS = "/device:CPU:0"
 
@@ -138,7 +138,7 @@ class TensorboardWriter(object):
                                   bucket_limit=bins.tolist(),
                                   bucket=counts.tolist())
 
-        def __create_histo_event(statistics_collector: BaseStatsCollector):
+        def __create_histo_event(statistics_collector: BaseStatsContainer):
             """
             Create an event of histogram, and attach it to a list of events outside
             the scope called 'events'.
@@ -186,7 +186,7 @@ class TensorboardWriter(object):
 
         """
 
-        def __get_node_attr(n: BaseNode) -> Dict[str, Any]:
+        def __get_node_attr(n: Node) -> Dict[str, Any]:
             """
             Create a dictionary to display as the node's attributes.
             The dictionary contains information from node's framework attributes, quantization attributes
@@ -203,10 +203,7 @@ class TensorboardWriter(object):
             if n.quantization_attr is not None:
                 attr.update(n.quantization_attr)
 
-            # To log quantization configurations we need to check
-            # if they exist at all, as we can log the initial graph,
-            # which its nodes do not have configurations yet.
-            # Log final config or unified candidates, not both
+            # log final config or unified candidates, not both
             if n.final_weights_quantization_cfg is not None:
                 attr.update(n.final_weights_quantization_cfg.__dict__)
             elif n.candidates_weights_quantization_cfg is not None:
@@ -216,7 +213,7 @@ class TensorboardWriter(object):
                 attr.update(n.activation_quantization_cfg.__dict__)
             return attr
 
-        def __get_node_output_dims(n: BaseNode) -> List[tuple]:
+        def __get_node_output_dims(n: Node) -> List[tuple]:
             """
             Get node's output shapes. If the first dimension in an output shape is None,
             it means the batch size is dynamic, and it's replaced with -1 to mark it.
@@ -240,7 +237,7 @@ class TensorboardWriter(object):
                 dims = [(-1,) + output_shape[1:] if output_shape[0] is None else output_shape]
             return dims
 
-        def __create_node_stats(n: BaseNode):
+        def __create_node_stats(n: Node):
             """
             Create a NodeExecStats for a node in the graph. A NodeExecStats contains the
             memory and compute time a node requires.

model_compression_toolkit/keras/back2framework/instance_builder.py

@@ -21,7 +21,7 @@ from networkx.algorithms.dag import topological_sort
 
 from tensorflow.keras.layers import Layer
 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.keras.constants import LAYER_NAME
 
 
@@ -35,7 +35,7 @@ class OperationHandler(object):
         self.node_to_fw_op_dict = instance_builder(self.node_sort)  # hold dictionary from node to its equivalent
         # Keras layer
 
-    def get_node_op_function(self, n: BaseNode) -> Layer:
+    def get_node_op_function(self, n: Node) -> Layer:
         """
         Get the Keras layer that was built from the passed node.
 
@@ -58,7 +58,7 @@ class OperationHandler(object):
         return op_func
 
 
-def node_builder(n: common.BaseNode) -> Layer:
+def node_builder(n: common.Node) -> Layer:
     """
     Build a Keras layer from a node.
 
@@ -78,7 +78,7 @@ def node_builder(n: common.BaseNode) -> Layer:
     return node_instance
 
 
-def instance_builder(toposort: List[BaseNode]) -> Dict[BaseNode, Layer]:
+def instance_builder(toposort: List[Node]) -> Dict[Node, Layer]:
     """
     Build a dictionary of nodes to their corresponding Keras
     layers, given a list of nodes.

model_compression_toolkit/keras/back2framework/model_builder.py

@@ -20,26 +20,24 @@ import tensorflow as tf
 if tf.__version__ < "2.6":
     from tensorflow.keras.layers import Input
     from tensorflow.python.keras.layers.core import TFOpLambda
-    from tensorflow.python.keras.engine.base_layer import TensorFlowOpLayer
-    from tensorflow.python.keras.layers import Layer
 else:
     from keras import Input
     from keras.layers.core import TFOpLambda
-    from keras.engine.base_layer import TensorFlowOpLayer, Layer
 
 from model_compression_toolkit.common.model_builder_mode import ModelBuilderMode
+from tensorflow.python.keras.engine.base_layer import TensorFlowOpLayer
+from tensorflow.python.keras.layers import Layer
 from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapper import QuantizeWrapper
 from typing import Tuple, Any, Dict, List
 from tensorflow.python.util.object_identity import Reference as TFReference
-from model_compression_toolkit.common.graph.functional_node import FunctionalNode
+
 
 from model_compression_toolkit import common
 from model_compression_toolkit.common.framework_info import FrameworkInfo
 from model_compression_toolkit.keras.default_framework_info import DEFAULT_KERAS_INFO
-from model_compression_toolkit.keras.quantizer.mixed_precision.quantization_config_factory import \
-    quantization_config_builder_mixed_precision
+from model_compression_toolkit.keras.quantizer.mixed_precision.quantization_config_factory import quantization_config_builder_mixed_precision
 from model_compression_toolkit.keras.quantizer.gradient_ptq.config_factory import quantization_config_builder_gptq
-from model_compression_toolkit.common import BaseNode, Graph
+from model_compression_toolkit.common import Node, Graph
 from model_compression_toolkit.common.graph.edge import EDGE_SINK_INDEX
 from model_compression_toolkit.keras.back2framework.instance_builder import OperationHandler
 from model_compression_toolkit.keras.reader.connectivity_handler import OutTensor
@@ -82,9 +80,9 @@ def is_layer_fake_quant(layer: Layer) -> bool:
             isinstance(layer, TFOpLambda) and layer.symbol == FQ_NODE_OP_V2_4)
 
 
-def build_input_tensors_list(node: BaseNode,
+def build_input_tensors_list(node: Node,
                              graph: Graph,
-                             node_to_output_tensors_dict: Dict[BaseNode, List[TFReference]]) -> List[List[TFReference]]:
+                             node_to_output_tensors_dict: Dict[Node, List[TFReference]]) -> List[List[TFReference]]:
     """
     Given a node, build a list of input tensors the node gets. The list is built
     based on the node's incoming edges and previous nodes' output tensors.
@@ -107,10 +105,10 @@ def build_input_tensors_list(node: BaseNode,
     return input_tensors
 
 
-def run_operation(n: BaseNode,
+def run_operation(n: Node,
                   input_tensors: List[List[TFReference]],
                   op_func: Layer,
-                  input_nodes_to_input_tensors: Dict[BaseNode, Any],
+                  input_nodes_to_input_tensors: Dict[Node, Any],
                   mode: ModelBuilderMode = ModelBuilderMode.QUANTIZED) -> List[TFReference]:
     """
     Applying the layer (op_func) to the input tensors (input_tensors).
@@ -131,43 +129,33 @@ def run_operation(n: BaseNode,
 
     if len(input_tensors) == 0:  # Placeholder handling
         out_tensors_of_n = input_nodes_to_input_tensors[n]
-        if n.is_activation_quantization_enabled():
-            if mode in [ModelBuilderMode.QUANTIZED, ModelBuilderMode.GPTQ, ModelBuilderMode.MIXEDPRECISION]:
-                # Adding a fake quant node to Input when in GPTQ mode because quantize_model doesn't quantize the
-                # input layer
-                fake_quant = n.activation_quantization_cfg.generate_quantization_node()
-
-                if fake_quant is None:
-                    raise Exception(f'{n.name} should be quantized, but activation quantization function is None')
-
+        if mode in [ModelBuilderMode.QUANTIZED, ModelBuilderMode.GPTQ, ModelBuilderMode.MIXEDPRECISION]:
+            # Adding a fake quant node to Input when in GPTQ mode because quantize_model doesn't quantize the input layer
+            assert n.activation_quantization_cfg is not None  # Input layers should always have activation config
+            fake_quant = n.activation_quantization_cfg.activation_quantization_fn(n.activation_quantization_cfg.activation_n_bits,
+                                                                       n.activation_quantization_cfg.activation_is_signed,
+                                                                       n.activation_quantization_cfg.activation_quantization_params)
+            if fake_quant is not None:
                 out_tensors_of_n = fake_quant(out_tensors_of_n)
-
+            
     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
-        if isinstance(n, FunctionalNode):
-            if n.inputs_as_list:  # If the first argument should be a list of tensors:
-                out_tensors_of_n = op_func(input_tensors, *n.op_call_args, **n.op_call_kwargs)
-            else:  # If the input tensors should not be a list but iterated:
-                out_tensors_of_n = op_func(*input_tensors, *n.op_call_args, **n.op_call_kwargs)
+        
+        # If operator expects a single input tensor, it cannot be a list as it should
+        # have a dtype field.
+        if len(input_tensors) == 1:
+            out_tensors_of_n = op_func(input_tensors[0], **n.op_call_args)
         else:
-            # If operator expects a single input tensor, it cannot be a list as it should
-            # have a dtype field.
-            if len(input_tensors) == 1:
-                input_tensors = input_tensors[0]
-            out_tensors_of_n = op_func(input_tensors)
+            out_tensors_of_n = op_func(input_tensors, **n.op_call_args)
 
         # Add a fake quant node if the node has an activation threshold.
-        if n.is_activation_quantization_enabled():
-            if mode in [ModelBuilderMode.QUANTIZED,
-                        ModelBuilderMode.MIXEDPRECISION]:
-
-                fake_quant = n.activation_quantization_cfg.generate_quantization_node()
-
-                if fake_quant is None:
-                    raise Exception(f'{n.name} should be quantized, but activation quantization function is None')
-
-                out_tensors_of_n = fake_quant(out_tensors_of_n)
+        if n.activation_quantization_cfg is not None:
+            if mode in [ModelBuilderMode.QUANTIZED, ModelBuilderMode.MIXEDPRECISION] and n.activation_quantization_cfg.enable_activation_quantization:
+                fake_quant = n.activation_quantization_cfg.activation_quantization_fn(n.activation_quantization_cfg.activation_n_bits,
+                                                                           n.activation_quantization_cfg.activation_is_signed,
+                                                                           n.activation_quantization_cfg.activation_quantization_params)
+                if fake_quant is not None:
+                    out_tensors_of_n = fake_quant(out_tensors_of_n)
 
     return out_tensors_of_n
 
@@ -277,11 +265,10 @@ def model_builder(graph: common.Graph,
             nodes = graph.find_node_by_name(get_node_name_from_layer(layer))
             if len(nodes) == 1:
                 node = nodes[0]
-                # Wrap only if its weights should be quantized
-                if node.is_weights_quantization_enabled():
-                    return QuantizeWrapper(layer, quantization_config_builder_mixed_precision(node, fw_info))
-                return layer
-
+                # does not need to get wrapped as its weights are not quantized
+                if node.candidates_weights_quantization_cfg is None:
+                    return layer
+                return QuantizeWrapper(layer, quantization_config_builder_mixed_precision(node, fw_info))
             elif is_layer_fake_quant(layer):
                 return layer
             else:

model_compression_toolkit/keras/constants.py

@@ -38,9 +38,6 @@ PAD_SAME = 'same'
 RELU_MAX_VALUE = 'max_value'
 THRESHOLD = 'threshold'
 NEGATIVE_SLOPE = 'negative_slope'
-CHANNELS_FORMAT = 'data_format'
-CHANNELS_FORMAT_FIRST = 'channels_first'
-CHANNELS_FORMAT_LAST = 'channels_last'
 
 # Layers variables names:
 KERNEL = 'kernel'

model_compression_toolkit/keras/default_framework_info.py

@@ -14,18 +14,12 @@
 # ==============================================================================
 
 
-import tensorflow as tf
-if tf.__version__ < "2.6":
-    from tensorflow.keras.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Reshape, ZeroPadding2D, Dropout, \
-        MaxPooling2D, Activation, ReLU, GlobalAveragePooling2D, Add, Multiply, AveragePooling2D, UpSampling2D, InputLayer, \
-        Concatenate, Softmax, PReLU, Flatten, Cropping2D
-else:
-    from keras.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Reshape, ZeroPadding2D, \
-        Dropout, MaxPooling2D, Activation, ReLU, GlobalAveragePooling2D, Add, Multiply, AveragePooling2D, UpSampling2D, \
-        InputLayer, Concatenate, Softmax, PReLU, Flatten, Cropping2D
+from tensorflow.keras.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Reshape, ZeroPadding2D, Dropout, \
+    MaxPooling2D, Activation, ReLU, GlobalAveragePooling2D, Add, Multiply, AveragePooling2D, UpSampling2D, InputLayer, \
+    Concatenate, Softmax, PReLU, Flatten, Cropping2D
 
 from model_compression_toolkit.common.defaultdict import DefaultDict
-from model_compression_toolkit.common.framework_info import FrameworkInfo, ChannelAxis
+from model_compression_toolkit.common.framework_info import FrameworkInfo
 from model_compression_toolkit.common.quantization.quantization_config import QuantizationMethod
 from model_compression_toolkit.common.quantization.quantizers.kmeans_quantizer import kmeans_quantizer
 from model_compression_toolkit.common.quantization.quantizers.lut_kmeans_quantizer import lut_kmeans_quantizer
@@ -51,9 +45,7 @@ NO_QUANTIZATION = [Reshape,
                    Cropping2D,
                    ZeroPadding2D,
                    Dropout,
-                   MaxPooling2D,
-                   tf.split,
-                   tf.quantization.fake_quant_with_min_max_vars]  # TODO:  replace with marking
+                   MaxPooling2D] # TODO:  replace with marking
 
 ACTIVATION = [Activation,
               ReLU,
@@ -65,19 +57,7 @@ ACTIVATION = [Activation,
               UpSampling2D,
               InputLayer,
               Concatenate,
-              PReLU,
-              tf.add,
-              tf.multiply,
-              tf.reduce_mean,
-              tf.reduce_min,
-              tf.reduce_sum,
-              tf.reduce_max,
-              tf.image.resize,
-              tf.image.crop_and_resize,
-              tf.concat]
-
-
-
+              PReLU]
 
 """
 Map each layer to a list of its' weights attributes that should get quantized.
@@ -131,11 +111,6 @@ WEIGHTS_QUANTIZER_MAPPING = {QuantizationMethod.POWER_OF_TWO: power_of_two_quant
                              QuantizationMethod.KMEANS: kmeans_quantizer,
                              QuantizationMethod.LUT_QUANTIZER: lut_kmeans_quantizer}
 
-"""
-Output channel index of the model's layers
-"""
-OUTPUT_CHANNEL_INDEX = ChannelAxis.NHWC
-
 DEFAULT_KERAS_INFO = FrameworkInfo(KERNEL_OPS,
                                    ACTIVATION,
                                    NO_QUANTIZATION,
@@ -144,5 +119,4 @@ DEFAULT_KERAS_INFO = FrameworkInfo(KERNEL_OPS,
                                    DEFAULT_CHANNEL_AXIS_DICT,
                                    ACTIVATION2MINMAX,
                                    LAYER2MINMAX,
-                                   KERNEL_ATTRIBUTES,
-                                   OUTPUT_CHANNEL_INDEX)
+                                   KERNEL_ATTRIBUTES)

model_compression_toolkit/keras/gradient_ptq/graph_info.py

@@ -20,14 +20,14 @@ from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapp
 from typing import Tuple, List
 
 from model_compression_toolkit.common.graph.base_graph import Graph
-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 USE_BIAS
 from model_compression_toolkit.keras.quantizer.gradient_ptq import ActivationQuantizeConfig, WeightQuantizeConfig, ActivationAndWeightQuantizeConfig
 from model_compression_toolkit.common.framework_info import FrameworkInfo
 from tensorflow.keras.models import Model
 
 
-def get_compare_points(input_graph: Graph) -> Tuple[List[BaseNode], List[str]]:
+def get_compare_points(input_graph: Graph) -> Tuple[List[Node], List[str]]:
     """
     Create a list of nodes with weights in a graph and a corresponding list
     of their names for tensors comparison purposes.

model_compression_toolkit/keras/gradient_ptq/graph_update.py

@@ -15,14 +15,8 @@
 
 
 import copy
-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.engine.base_layer import TensorFlowOpLayer
-else:
-    from keras.engine.base_layer import TensorFlowOpLayer
 
+from tensorflow.python.keras.engine.base_layer import TensorFlowOpLayer
 from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapper import QuantizeWrapper
 
 from model_compression_toolkit import common

model_compression_toolkit/keras/gradient_ptq/training_wrapper.py

@@ -57,6 +57,7 @@ def gptq_training_wrapper(tg: Graph,
     #########################################
     # Build two models and compare points
     #########################################
+    # TODO: maybe need to add pre_build substitutions here. Ask Elad
     compare_points, _ = get_compare_points(tg)  # get compare points
     n = len(compare_points)
     float_model, float_user_info = model_builder(tg,

model_compression_toolkit/keras/graph_substitutions/substitutions/activation_decomposition.py

@@ -21,7 +21,7 @@ from model_compression_toolkit.common.constants import FLOAT_32, DATA_TYPE
 from model_compression_toolkit.common.graph.base_graph import Graph
 from model_compression_toolkit.common.graph.graph_matchers import NodeOperationMatcher, \
     NodeFrameworkAttrMatcher
-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 LINEAR, ACTIVATION, TRAINABLE, LAYER_NAME
 
 
@@ -49,7 +49,7 @@ class ActivationDecomposition(common.BaseSubstitution):
 
     def substitute(self,
                    graph: Graph,
-                   op2d_node: BaseNode) -> Graph:
+                   op2d_node: Node) -> Graph:
         """
         Decompose the activation function in a linear node to a new activation layer.
         Set activation function in the linear node to 'linear' (y=x).
@@ -70,14 +70,12 @@ class ActivationDecomposition(common.BaseSubstitution):
             DATA_TYPE: FLOAT_32,
             ACTIVATION: op2d_node.framework_attr.get(ACTIVATION)}
 
-        activation_node = common.graph.BaseNode(activation_node_name,
-                                                activation_fw_attr,
-                                                op2d_node.output_shape,
-                                                op2d_node.output_shape,
-                                                {},
-                                                Activation)
-
-
+        activation_node = common.graph.Node(activation_node_name,
+                                            activation_fw_attr,
+                                            op2d_node.output_shape,
+                                            op2d_node.output_shape,
+                                            {},
+                                            Activation)
 
         graph.add_node(activation_node)
         graph.reconnect_out_edges(current_node=op2d_node,

model_compression_toolkit/keras/graph_substitutions/substitutions/batchnorm_folding.py

@@ -24,7 +24,7 @@ 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 EdgeMatcher, NodeOperationMatcher, \
     NodeFrameworkAttrMatcher
-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, BIAS, USE_BIAS, LINEAR, ACTIVATION, LAYER_NAME, \
     GAMMA, BETA, EPSILON, \
     MOVING_MEAN, \
@@ -51,7 +51,7 @@ class BatchNormalizationFolding(common.BaseSubstitution):
 
     def substitute(self,
                    graph: Graph,
-                   edge_nodes: Tuple[BaseNode, BaseNode]) -> Graph:
+                   edge_nodes: Tuple[Node, Node]) -> Graph:
         """
         Fold BatchNormalization into preceding linear layers.
 

model_compression_toolkit/keras/graph_substitutions/substitutions/input_scaling.py

@@ -21,7 +21,7 @@ from model_compression_toolkit import common
 from model_compression_toolkit.common.framework_info import FrameworkInfo
 from model_compression_toolkit.common.graph.base_graph import Graph
 from model_compression_toolkit.common.graph.graph_matchers import NodeOperationMatcher, EdgeMatcher, WalkMatcher
-from model_compression_toolkit.common.graph.base_node import BaseNode
+from model_compression_toolkit.common.graph.node import Node
 from model_compression_toolkit.common.quantization.quantization_config import QuantizationConfig
 from model_compression_toolkit.common.constants import THRESHOLD
 from model_compression_toolkit.keras.constants import KERNEL
@@ -62,7 +62,7 @@ class BaseInputScaling(common.BaseSubstitution):
 
     def substitute(self,
                    graph: Graph,
-                   nodes_list: List[BaseNode]) -> Graph:
+                   nodes_list: List[Node]) -> Graph:
         """
         Scale activation threshold for input layers, if they are followed by linear nodes. We first
         scale their thresholds to a constrained threshold, and then fix it by scaling the linear op weights

model_compression_toolkit/keras/graph_substitutions/substitutions/mark_activation.py

@@ -21,7 +21,7 @@ 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, EdgeMatcher, \
     NodeFrameworkAttrMatcher
-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 LINEAR, ACTIVATION
 
 
@@ -55,7 +55,7 @@ class MarkActivation(common.BaseSubstitution):
 
     def substitute(self,
                    graph: Graph,
-                   edge: Tuple[BaseNode, BaseNode]) -> Graph:
+                   edge: Tuple[Node, Node]) -> Graph:
         """
         Mark the first node in an edge that should not be quantized as so.
         This can be done due to the following reasons:
@@ -69,5 +69,5 @@ class MarkActivation(common.BaseSubstitution):
             Graph after applying the substitution.
         """
 
-        edge[0].activation_quantization_cfg.enable_activation_quantization = False
+        edge[0].output_quantization = False
         return graph

model_compression_toolkit/keras/graph_substitutions/substitutions/relu_bound_correction.py

@@ -20,11 +20,12 @@ import numpy as np
 from tensorflow.keras.layers import ReLU, Activation, DepthwiseConv2D, Conv2DTranspose, Conv2D, Dense
 
 from model_compression_toolkit import common
-from model_compression_toolkit.common import Graph, BaseNode
+from model_compression_toolkit.common import FrameworkInfo, Graph, Node
 from model_compression_toolkit.common.constants import THRESHOLD
 from model_compression_toolkit.common.graph.graph_matchers import NodeOperationMatcher, WalkMatcher, \
     NodeFrameworkAttrMatcher
-from model_compression_toolkit.common.collectors.statistics_collector import is_number
+from model_compression_toolkit.common.quantization.quantization_config import QuantizationConfig
+from model_compression_toolkit.common.statistics_collector import is_number
 from model_compression_toolkit.keras.constants import KERNEL, BIAS, ACTIVATION, RELU_MAX_VALUE
 from model_compression_toolkit.keras.constants import RELU
 
@@ -60,7 +61,7 @@ class ReLUBoundCorrection(common.BaseSubstitution):
 
     def substitute(self,
                    graph: Graph,
-                   nodes_list: List[BaseNode]) -> Graph:
+                   nodes_list: List[Node]) -> Graph:
         """
         Transform a list of nodes in a graph to use the entire constrained quantized range.
         This is done by scaling two linear nodes with a non-linearity between them, if the non-linearity