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

model_compression_toolkit/keras/quantizer/gradient_ptq/config_factory.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 
-from tensorflow.keras.layers import DepthwiseConv2D
+from tensorflow.keras.layers import DepthwiseConv2D, Conv2D, Conv2DTranspose, Dense
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_quantize_configs import \
     NoOpQuantizeConfig
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_quantize_registry import \
@@ -30,7 +30,7 @@ from model_compression_toolkit.common.framework_info import FrameworkInfo
 MAX_LSBS_CHANGE = 8
 
 
-def quantization_config_builder_gptq(n: common.BaseNode,
+def quantization_config_builder_gptq(n: common.Node,
                                      fw_info: FrameworkInfo) -> QuantizeConfig:
     """
     Build a QuantizeConfig for a node according to its quantization configuration and
@@ -45,28 +45,29 @@ def quantization_config_builder_gptq(n: common.BaseNode,
         quantization configuration).
     """
 
-    if n.is_weights_quantization_enabled() and n.is_activation_quantization_enabled():
+    if n.activation_quantization() and n.weight_quantization():
         qc = keras.quantizer.gradient_ptq.ActivationAndWeightQuantizeConfig(fw_info.get_kernel_op_attributes(n.layer_class),
                                                                             n.final_weights_quantization_cfg.weights_quantization_params.get(THRESHOLD),
                                                                             n.final_weights_quantization_cfg.weights_channels_axis,
                                                                             n.final_weights_quantization_cfg.weights_n_bits,
                                                                             n.activation_quantization_cfg.activation_quantization_params,
+                                                                            n.activation_quantization_cfg.activation_is_signed,
                                                                             activation_num_bits=n.activation_quantization_cfg.activation_n_bits,
                                                                             max_lsbs_change=MAX_LSBS_CHANGE
                                                                             )
-
-    elif n.is_activation_quantization_enabled() and not n.is_weights_quantization_enabled():
+    elif n.activation_quantization():
         qc = keras.quantizer.gradient_ptq.ActivationQuantizeConfig(n.activation_quantization_cfg.activation_quantization_params,
+                                                                   n.activation_quantization_cfg.activation_is_signed,
                                                                    num_bits=n.activation_quantization_cfg.activation_n_bits)
-
-    elif n.is_weights_quantization_enabled() and not n.is_activation_quantization_enabled():
+    elif n.weight_quantization():
         qc = keras.quantizer.gradient_ptq.WeightQuantizeConfig(fw_info.get_kernel_op_attributes(n.layer_class),
                                                                n.final_weights_quantization_cfg.weights_quantization_params.get(THRESHOLD),
                                                                n.final_weights_quantization_cfg.weights_channels_axis,
                                                                n.final_weights_quantization_cfg.weights_n_bits,
-                                                               max_lsbs_change=MAX_LSBS_CHANGE)
+                                                               max_lsbs_change=MAX_LSBS_CHANGE
+                                                               )
 
-    elif not n.is_weights_quantization_enabled() and not n.is_activation_quantization_enabled():
+    elif n.no_quantization():
         qc = NoOpQuantizeConfig()
 
     else:

model_compression_toolkit/keras/quantizer/gradient_ptq/weight_quantizer_gptq_config.py

@@ -16,15 +16,7 @@
 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
 

model_compression_toolkit/keras/quantizer/mixed_precision/quantization_config_factory.py

@@ -25,7 +25,7 @@ from model_compression_toolkit.keras.quantizer.mixed_precision.selective_weights
 
 
 
-def quantization_config_builder_mixed_precision(n: common.BaseNode,
+def quantization_config_builder_mixed_precision(n: common.Node,
                                                 fw_info: FrameworkInfo) -> QuantizeConfig:
     """
     Build a QuantizeConfig for layers that should be wrapped in a QuantizeWrapper to

model_compression_toolkit/keras/quantizer/mixed_precision/selective_weights_quantize_config.py

@@ -17,12 +17,7 @@ from typing import List, Tuple, Any, Dict
 
 import numpy as np
 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.quantize_config import QuantizeConfig
 

model_compression_toolkit/keras/reader/common.py

@@ -20,19 +20,20 @@ import tensorflow as tf
 if tf.__version__ < "2.6":
     from tensorflow.python.keras.engine.node import Node as KerasNode
     from tensorflow.keras.layers import InputLayer
-    from tensorflow.python.keras.engine.functional import Functional
-    from tensorflow.python.keras.engine.sequential import Sequential
 else:
     from keras.engine.input_layer import InputLayer
     from keras.engine.node import Node as KerasNode
-    from keras.engine.functional import Functional
-    from keras.engine.sequential import Sequential
 
-from model_compression_toolkit.common.graph.base_node import BaseNode
 
+from tensorflow.python.keras.engine.functional import Functional
 
+from tensorflow.python.keras.engine.sequential import Sequential
 
-def is_node_an_input_layer(node: BaseNode) -> bool:
+from model_compression_toolkit.common.graph.node import Node
+
+
+
+def is_node_an_input_layer(node: Node) -> bool:
     """
     Checks if a node represents a Keras input layer.
     Args:
@@ -41,14 +42,14 @@ def is_node_an_input_layer(node: BaseNode) -> bool:
     Returns:
         Whether the node represents an input layer or not.
     """
-    if isinstance(node, BaseNode):
+    if isinstance(node, Node):
         return node.layer_class == InputLayer
     elif isinstance(node, KerasNode):
         return isinstance(node.layer, InputLayer)
     else:
         raise Exception('Node to check has to be either a graph node or a keras node')
 
-def is_node_a_model(node: BaseNode) -> bool:
+def is_node_a_model(node: Node) -> bool:
     """
     Checks if a node represents a Keras model.
     Args:
@@ -57,10 +58,11 @@ def is_node_a_model(node: BaseNode) -> bool:
     Returns:
         Whether the node represents a Keras model or not.
     """
-    if isinstance(node, BaseNode):
+    if isinstance(node, Node):
         return node.layer_class in [Functional, Sequential]
     elif isinstance(node, KerasNode):
         return isinstance(node.layer, Functional) or isinstance(node.layer, Sequential)
     else:
         raise Exception('Node to check has to be either a graph node or a keras node')
 
+    # return node.layer_class in [Functional, Sequential]

model_compression_toolkit/keras/reader/connectivity_handler.py

@@ -15,18 +15,12 @@
 
 
 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.node import Node as KerasNode
-else:
-    from keras.engine.node import Node as KerasNode
-
-
+from tensorflow.python.keras.engine.node import Node as KerasNode
 from tensorflow.python.util.object_identity import Reference as TFReference
 from typing import List, Tuple
 
 from model_compression_toolkit.common.graph.base_graph import OutTensor
-from model_compression_toolkit.common.graph.base_node import BaseNode
+from model_compression_toolkit.common.graph.node import Node
 from model_compression_toolkit.keras.reader.common import is_node_an_input_layer
 from model_compression_toolkit.keras.reader.node_builder import build_node
 
@@ -47,7 +41,7 @@ class ConnectivityHandler(object):
         self._nodes2output_tensors = dict()  # Node -> List[Tensor]
         self._output_tensors2nodes = dict()  # Tensor -> Node
 
-    def get_nodes(self) -> List[BaseNode]:
+    def get_nodes(self) -> List[Node]:
         """
         Returns: List of nodes in the connectivity handler.
         """
@@ -66,7 +60,7 @@ class ConnectivityHandler(object):
         return self._input_tensors2nodes.get(tensor) is not None
 
     def input_tensor2nodes(self,
-                           in_tensor: str) -> List[BaseNode]:
+                           in_tensor: str) -> List[Node]:
         """
         Returns a list of nodes that have a given tensor in their input tensors.
         Args:
@@ -90,7 +84,7 @@ class ConnectivityHandler(object):
         return self._output_tensors2nodes[out_tensor] if out_tensor in self._output_tensors2nodes else None
 
     def node2input_tensors(self,
-                           node: BaseNode) -> List[TFReference]:
+                           node: Node) -> List[TFReference]:
         """
         Get a list of input tensors of a node.
         Args:
@@ -102,7 +96,7 @@ class ConnectivityHandler(object):
         return self._nodes2input_tensors[node] if node in self._nodes2input_tensors else []
 
     def node2output_tensors(self,
-                            node: BaseNode) -> List[TFReference]:
+                            node: Node) -> List[TFReference]:
         """
         Get a list of output tensors of a node.
         Args:
@@ -179,8 +173,8 @@ class ConnectivityHandler(object):
             self._output_tensors2nodes[output_t] = node
 
     def get_edge_indices(self,
-                         src_node: BaseNode,
-                         dst_node: BaseNode,
+                         src_node: Node,
+                         dst_node: Node,
                          connecting_tensor: TFReference) -> Tuple[int, int]:
         """
         Get indices of an edge by its source/destination nodes and the connecting tensor which defines the edge.
@@ -199,7 +193,7 @@ class ConnectivityHandler(object):
         return src_index, dst_index
 
     def get_out_edges_params_list(self,
-                                  src_node: BaseNode) -> List[tuple]:
+                                  src_node: Node) -> List[tuple]:
         """
         Compute for a given node, all parameters of its outgoing edges.
         Args:

model_compression_toolkit/keras/reader/nested_model/edges_merger.py

@@ -18,11 +18,11 @@ from typing import List, Dict, Tuple
 
 from model_compression_toolkit.common.graph.base_graph import Graph
 from model_compression_toolkit.common.graph.edge import Edge, convert_to_edge
-from model_compression_toolkit.common.graph.base_node import BaseNode
+from model_compression_toolkit.common.graph.node import Node
 from model_compression_toolkit.keras.reader.connectivity_handler import OutTensor
 
 
-def merge_models_edges(inner_model_node: BaseNode,
+def merge_models_edges(inner_model_node: Node,
                        outer_graph: Graph,
                        inner_graph: Graph) -> List[Edge]:
     """
@@ -91,8 +91,8 @@ def rewire_outgoing_edged(inner_model_outputs: List[OutTensor],
                 res_edges.remove(model_node_out_edge)
 
 
-def rewire_incoming_edges(inner_inputs_out_edges: Dict[BaseNode, List[Edge]],
-                          inner_model_inputs_dict: List[BaseNode],
+def rewire_incoming_edges(inner_inputs_out_edges: Dict[Node, List[Edge]],
+                          inner_model_inputs_dict: List[Node],
                           model_node_in_edges: List[Edge],
                           res_edge_list: List[Edge]):
     """
@@ -126,7 +126,7 @@ def rewire_incoming_edges(inner_inputs_out_edges: Dict[BaseNode, List[Edge]],
         res_edge_list.remove(model_node_in_edge)
 
 
-def get_inner_inputs_successors(inner_graph: Graph) -> Dict[BaseNode, List[Edge]]:
+def get_inner_inputs_successors(inner_graph: Graph) -> Dict[Node, List[Edge]]:
     """
     Compute out edges the input nodes of the inner model has.
     Args:
@@ -147,7 +147,7 @@ def get_inner_inputs_successors(inner_graph: Graph) -> Dict[BaseNode, List[Edge]
     return inner_inputs_out_edges
 
 
-def get_model_node_edges(model_node: BaseNode,
+def get_model_node_edges(model_node: Node,
                          outer_edge_list: List[Edge]) -> Tuple[List[Edge], List[Edge]]:
     """
     Get incoming and outgoing edges the inner model node has in the outer graph.

model_compression_toolkit/keras/reader/nested_model/nested_model_handler.py

@@ -15,13 +15,13 @@
 
 
 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.reader.nested_model.edges_merger import merge_models_edges
 from model_compression_toolkit.keras.reader.nested_model.nodes_merger import merge_models_nodes
 from model_compression_toolkit.keras.reader.nested_model.outputs_merger import merge_models_outputs
 
 
-def merge_graphs(inner_model_node: BaseNode,
+def merge_graphs(inner_model_node: Node,
                  outer_graph: Graph,
                  inner_graph: Graph) -> Graph:
     """

model_compression_toolkit/keras/reader/nested_model/nodes_merger.py

@@ -20,12 +20,12 @@ import copy
 from typing import 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
 
 
-def merge_models_nodes(inner_model_node: BaseNode,
+def merge_models_nodes(inner_model_node: Node,
                        outer_graph: Graph,
-                       inner_graph: Graph) -> List[BaseNode]:
+                       inner_graph: Graph) -> List[Node]:
     """
     Given two MultiDiGraphs (one of an outer model and the second of the inner model), merge their nodes into
     a single nodes list representing the nodes that should be in a single MultiDiGraph after unrolling the inner graph.

model_compression_toolkit/keras/reader/nested_model/outputs_merger.py

@@ -19,11 +19,11 @@ import copy
 from typing import 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.reader.connectivity_handler import OutTensor
 
 
-def merge_models_outputs(inner_model_node: BaseNode,
+def merge_models_outputs(inner_model_node: Node,
                          outer_graph: Graph,
                          inner_graph: Graph) -> List[OutTensor]:
     """

model_compression_toolkit/keras/reader/node_builder.py

@@ -12,21 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Any
 
-import tensorflow as tf
 
-if tf.__version__ < "2.6":
-    from tensorflow.python.keras.layers.core import TFOpLambda, SlicingOpLambda
-    from tensorflow.python.keras.engine.keras_tensor import KerasTensor
-    from tensorflow.python.keras.engine.node import Node as KerasNode
-else:
-    from keras.layers.core import TFOpLambda, SlicingOpLambda
-    from keras.engine.keras_tensor import KerasTensor
-    from keras.engine.node import Node as KerasNode
+import tensorflow as tf
+from tensorflow.python.keras.engine.node import Node as KerasNode
 
-from model_compression_toolkit.common.graph.base_node import BaseNode
-from model_compression_toolkit.common.graph.functional_node import FunctionalNode
+from model_compression_toolkit.common.graph.node import Node
 
 keras = tf.keras
 layers = keras.layers
@@ -35,7 +26,7 @@ REUSED_IDENTIFIER = '_reused_'
 
 
 def build_node(node: KerasNode,
-               node_name_to_node: dict) -> BaseNode:
+               node_name_to_node: dict) -> Node:
     """
     Build a node from a Keras node. A node contains all information to reconstruct the layer it's representing
     in a model:
@@ -51,8 +42,7 @@ def build_node(node: KerasNode,
     """
     keras_layer = node.layer  # get the layer the node represents.
     layer_config = keras_layer.get_config()  # layer configuration to reconstruct it.
-    op_call_args = node.call_args
-    op_call_kwargs = node.call_kwargs
+    op_call_args = node.call_kwargs
     layer_class = type(keras_layer)  # class path to instantiating it in back2framework.
     weights = {v.name: v.numpy() for v in keras_layer.weights}  # layer's weights
 
@@ -76,56 +66,16 @@ def build_node(node: KerasNode,
     input_shape = keras_layer.get_input_shape_at(io_index)
     output_shape = keras_layer.get_output_shape_at(io_index)
 
-    if layer_class in [TFOpLambda, SlicingOpLambda]:
-        # Some functional ops (such as tf.concat) should receive the input tensors as a list
-        # and some are not (such as tf.multiply), so each FunctionalNode holds
-        # a flag to indicate that.
-        inputs_as_list = __is_functional_inputs_a_list(op_call_args)
-        # Do not hold the tensors that are in op_call_args as they are
-        # not needed. Thus, if the first element in op_call_args is a list of
-        # Keras tensors, remove it from op_call_args.
-        op_call_args = op_call_args[int(inputs_as_list):]
-        node = FunctionalNode(node_name,
-                              layer_config,
-                              input_shape,
-                              output_shape,
-                              weights,
-                              layer_class,
-                              [arg for arg in op_call_args if not isinstance(arg, KerasTensor)],  # Do not hold the tensors that are in op_call_args
-                              {k: v for k, v in op_call_kwargs.items() if not isinstance(v, KerasTensor)}, # In TF2.5 tensors are in kwargs as well.
-                              is_reused,
-                              reuse_group,
-                              functional_op=keras_layer.function,
-                              inputs_as_list=inputs_as_list)
-    else:
-        node = BaseNode(node_name,
-                        layer_config,
-                        input_shape,
-                        output_shape,
-                        weights,
-                        layer_class,
-                        is_reused,
-                        reuse_group)
+    node = Node(node_name,
+                layer_config,
+                input_shape,
+                output_shape,
+                weights,
+                layer_class,
+                is_reused,
+                reuse_group,
+                op_call_args)
 
     node_name_to_node[node_name] = node
-    return node
-
-
-def __is_functional_inputs_a_list(op_call_args: Any) -> bool:
-    """
-    Check whether the input tensors should be passed as a list
-    or not.
 
-    Args:
-        op_call_args: Arguments list to check.
-
-    Returns:
-        Whether the input tensors should be passed as a list or not.
-    """
-
-    if len(op_call_args) > 0 and isinstance(op_call_args[0], list):
-        inputs_as_list = True
-        for arg in op_call_args[0]:
-            inputs_as_list = inputs_as_list and isinstance(arg, KerasTensor)
-        return inputs_as_list
-    return False
+    return node

model_compression_toolkit/keras/reader/reader.py

@@ -24,8 +24,8 @@ from typing import List
 
 from model_compression_toolkit.common.graph.base_graph import Graph
 from model_compression_toolkit.common.graph.edge import Edge
-from model_compression_toolkit.common.graph.base_node import BaseNode
-from model_compression_toolkit.keras.reader.common import is_node_a_model, is_node_an_input_layer
+from model_compression_toolkit.common.graph.node import Node
+from model_compression_toolkit.keras.reader.common import is_node_a_model
 from model_compression_toolkit.keras.reader.connectivity_handler import ConnectivityHandler
 from model_compression_toolkit.keras.reader.nested_model.nested_model_handler import merge_graphs
 
@@ -70,6 +70,7 @@ def build_tensors_list(tensors_list) -> List[TFReference]:
                                                                                      tensor in
                                                                                      tensors_list]
 
+
 def build_connectivity_handler(model: Model) -> ConnectivityHandler:
     """
     Build a connectivity handler containing all information about connections in the model (nodes and
@@ -84,8 +85,7 @@ def build_connectivity_handler(model: Model) -> ConnectivityHandler:
     connectivity_handler = ConnectivityHandler()
     for nodes in model._nodes_by_depth.values():
         for node in nodes:  # nodes_by depth values are lists (each list for a different depth)
-            node_inputs = node.input_tensors if is_node_an_input_layer(node) else node.keras_inputs
-            input_tensors = build_tensors_list(node_inputs)  # build input tensors of the node
+            input_tensors = build_tensors_list(node.input_tensors)  # build input tensors of the node
             output_tensors = build_tensors_list(node.output_tensors)  # build output tensors of the node
             connectivity_handler.add_node(node,
                                           input_tensors,
@@ -143,7 +143,7 @@ def parse_model(model: Model) -> Graph:
 
 
 def flatten_nested_model(outer_graph: Graph,
-                         inner_model_node: BaseNode,
+                         inner_model_node: Node,
                          outer_keras_model: Model):
     """
     Flat a nested model given two graphs: inner and outer models' graphs.

model_compression_toolkit/keras/tensor_marking.py

@@ -0,0 +1,113 @@
+# Copyright 2021 Sony Semiconductors Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+from tensorflow.keras.layers import ReLU, Activation
+
+from model_compression_toolkit import common
+from model_compression_toolkit.common import FrameworkInfo
+from model_compression_toolkit.common.graph.node import Node
+from model_compression_toolkit.common.statistics_collector import BaseStatsContainer
+from model_compression_toolkit.keras.constants import LINEAR, ACTIVATION, RELU_MAX_VALUE, THRESHOLD, NEGATIVE_SLOPE
+
+
+def get_stats_collector_for_activation_op(n: Node,
+                                          fw_info: FrameworkInfo) -> common.StatsContainer:
+    """
+    Create and initial a statistics collector for an activation layer. If the activation function's min/max
+    output values are known, the statistics collector is initialized with these values.
+
+    Args:
+        fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
+        groups of layers by how they should be quantized, etc.)
+        n: Node to create a statistics collector for it.
+
+    Returns:
+        Statistics collector initialized with known min/max values.
+    """
+
+    if n.layer_class == ReLU:
+        negative_slope = n.framework_attr[NEGATIVE_SLOPE]
+        threshold = n.framework_attr[THRESHOLD]
+        init_max = n.framework_attr[RELU_MAX_VALUE]
+        return common.StatsContainer(init_min_value=threshold if negative_slope == 0 else None,
+                                     init_max_value=init_max)
+
+    if n.layer_class == Activation:
+        init_min, init_max = fw_info.activation_min_max_mapping[n.framework_attr[ACTIVATION]]
+        return common.StatsContainer(init_min_value=init_min,
+                                     init_max_value=init_max)
+
+    if fw_info.layers_has_min_max(n.layer_class):
+        init_min, init_max = fw_info.layer_min_max_mapping[n.layer_class]
+        return common.StatsContainer(init_min_value=init_min,
+                                     init_max_value=init_max)
+
+    return common.StatsContainer()
+
+
+def get_stats_collector_for_kernel_op(n: common.Node,
+                                      fw_info: FrameworkInfo) -> BaseStatsContainer:
+    """
+    Create and initial a statistics collector for a linear operator. If the layer has an activation function and
+    its min/max output values are known, the statistics collector is initialized with these values.
+    If the layer's output should not be quantized, NoStatsContainer is created.
+
+    Args:
+        fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
+        groups of layers by how they should be quantized, etc.)
+        n: Node to create a statistics collector for it.
+
+    Returns:
+        BaseStatsContainer according to statistics that are collected.
+    """
+
+    if n.framework_attr[ACTIVATION] == LINEAR and n.output_quantization:
+        return common.StatsContainer()
+
+    if n.framework_attr[ACTIVATION] == LINEAR and not n.output_quantization:
+        return common.NoStatsContainer()
+
+    if n.framework_attr[ACTIVATION] in fw_info.activation_min_max_mapping.keys():
+        min_value, max_value = fw_info.activation_min_max_mapping[n.framework_attr[ACTIVATION]]
+        return common.StatsContainer(init_min_value=min_value,
+                                     init_max_value=max_value)
+
+    return common.StatsContainer()
+
+
+def get_node_stats_collector(node: common.Node,
+                             fw_info: common.FrameworkInfo) -> common.statistics_collector.BaseStatsContainer:
+    """
+    Gets a node and a groups list and create and return a statistics collector for the node
+    according to the group the node is in.
+
+    Args:
+        node: Node to create its statistics collector.
+        fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
+        groups of layers by how they should be quantized, etc.)
+
+    Returns:
+        Statistics collector for statistics collection for the node.
+    """
+
+    stats_collector = get_stats_collector_for_activation_op(node, fw_info)
+    if fw_info.in_no_quantization_ops(node):  # node should not be quantized
+        stats_collector = common.NoStatsContainer()
+
+    if fw_info.in_kernel_ops(node):  # node's kernel should be quantized
+        stats_collector = get_stats_collector_for_kernel_op(node, fw_info)
+
+    return stats_collector

model_compression_toolkit/keras/visualization/nn_visualizer.py

@@ -22,10 +22,10 @@ from model_compression_toolkit.common import Graph
 from model_compression_toolkit.common.similarity_analyzer import compute_cs
 from model_compression_toolkit.keras.back2framework.model_builder import model_builder
 from model_compression_toolkit.common.model_builder_mode import ModelBuilderMode
-from model_compression_toolkit.common.graph.base_node import BaseNode
+from model_compression_toolkit.common.graph.node import Node
 
 
-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 in a graph where we collect their output statistics, and a corresponding list
     of their names for tensors comparison purposes.

model_compression_toolkit/common/collectors/statistics_collector_generator.py

@@ -1,43 +0,0 @@
-# Copyright 2021 Sony Semiconductors Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-
-from model_compression_toolkit import common
-from model_compression_toolkit.common.collectors.statistics_collector import BaseStatsCollector
-
-
-def create_stats_collector_for_node(node: common.BaseNode,
-                                    output_channel_index: int) -> BaseStatsCollector:
-    """
-    Gets a node and a groups list and create and return a statistics collector for a node
-    according to whether its statistics should be collected and the prior information we
-    have about this node.
-
-    Args:
-        node: Node to create its statistics collector.
-        output_channel_index: Index of output channels (for statistics per-channel).
-
-    Returns:
-        Statistics collector for statistics collection for the node.
-    """
-
-    if node.is_activation_quantization_enabled():
-        stats_collector = common.StatsCollector(init_min_value=node.prior_info.min_output,
-                                                init_max_value=node.prior_info.max_output,
-                                                output_channel_index=output_channel_index)
-    else:
-        stats_collector = common.NoStatsCollector()
-
-    return stats_collector