mct-nightly 2.0.0.20240418.439__py3-none-any.whl → 2.0.0.20240419.358__py3-none-any.whl

model_compression_toolkit/core/keras/keras_implementation.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 from functools import partial
-from typing import List, Any, Tuple, Callable, Dict
+from typing import List, Any, Tuple, Callable, Dict, Union
 
 import numpy as np
 import tensorflow as tf
@@ -412,12 +412,13 @@ class KerasImplementation(FrameworkImplementation):
         Returns: True if the node should be considered an interest point, False otherwise.
         """
 
-        if node.type == Activation:
+        if node.is_match_type(Activation):
             node_type_name = node.framework_attr[keras_constants.ACTIVATION]
             if node_type_name in [keras_constants.SOFTMAX, keras_constants.SIGMOID]:
                 return True
-        elif node.type in [tf.nn.softmax, tf.keras.layers.Softmax, tf.nn.sigmoid, Conv2D, DepthwiseConv2D, Conv2DTranspose, Dense, Concatenate,
-                           tf.concat, Add, tf.add]:
+        elif any([node.is_match_type(_type) for _type in [tf.nn.softmax, tf.keras.layers.Softmax, tf.nn.sigmoid, Conv2D,
+                                                          DepthwiseConv2D, Conv2DTranspose, Dense, Concatenate, tf.concat,
+                                                          Add, tf.add]]):
             return True
 
         return False
@@ -529,18 +530,18 @@ class KerasImplementation(FrameworkImplementation):
         kernel_shape = node.get_weights_by_keys(fw_info.get_kernel_op_attributes(node.type)[0]).shape
         output_channel_axis, input_channel_axis = fw_info.kernel_channels_mapping.get(node.type)
 
-        if node.type is Conv2D or node.type is Conv2DTranspose:
+        if node.is_match_type(Conv2D) or node.is_match_type(Conv2DTranspose):
             # (C_out * W_out * H_out) * C_in * (W_kernel * H_kernel)
             return np.prod([x for x in output_shape if x is not None]) * \
                    kernel_shape[input_channel_axis] * \
                    (kernel_shape[0] * kernel_shape[1])
-        elif node.type is DepthwiseConv2D:
+        elif node.is_match_type(DepthwiseConv2D):
             # Depth * (W_out * H_out) * C_in * (W_kernel * H_kernel)
             return node.framework_attr.get(DEPTH_MULTIPLIER) * \
                    np.prod([x for x in output_shape if x is not None]) / output_shape[output_channel_axis] * \
                    kernel_shape[input_channel_axis] * \
                    (kernel_shape[0] * kernel_shape[1])
-        elif node.type is Dense:
+        elif node.is_match_type(Dense):
             # IN * OUT
             return kernel_shape[0] * kernel_shape[1]
         else:
@@ -593,10 +594,9 @@ class KerasImplementation(FrameworkImplementation):
         Returns:
             weight_quantizers: A dictionary between a weight's name to its quantizer.
             activation_quantizers: A list of activations quantization, one for each layer output.
-
         """
 
-        def _weight_name(w: str) -> str:
+        def _weight_name(w: Union[str, int]) -> Union[str, int]:
             """
             Extracts the weight name from the full TensorFlow variable name.
 
@@ -609,7 +609,7 @@ class KerasImplementation(FrameworkImplementation):
               Extracted weight name.
             """
 
-            return w.split(':')[0].split('/')[-1]
+            return w.split(':')[0].split('/')[-1] if isinstance(w, str) else w
 
         attribute_names = [_weight_name(wn) for wn in node.get_node_weights_attributes()
                            if node.is_weights_quantization_enabled(wn)]

model_compression_toolkit/core/keras/keras_node_prior_info.py

@@ -56,13 +56,13 @@ def _get_min_max_outputs(node: BaseNode,
     """
     min_output, max_output = None, None
 
-    if node.type == ReLU:
+    if node.is_match_type(ReLU):
         min_output = node.framework_attr[THRESHOLD] if node.framework_attr[NEGATIVE_SLOPE] == 0 else None
 
     elif fw_info.layers_has_min_max(node.type):
         min_output, max_output = fw_info.layer_min_max_mapping[node.type]
 
-    elif node.type == Activation and fw_info.activation_has_min_max(node.framework_attr[ACTIVATION]):
+    elif node.is_match_type(Activation) and fw_info.activation_has_min_max(node.framework_attr[ACTIVATION]):
         min_output, max_output = fw_info.activation_min_max_mapping[node.framework_attr[ACTIVATION]]
 
     return min_output, max_output
@@ -82,7 +82,7 @@ def _get_mean_std_outputs(node: BaseNode,
     """
     mean_output, std_output = None, None
 
-    if node.type == BatchNormalization:
+    if node.is_match_type(BatchNormalization):
         mean_output = node.get_weights_by_keys(BETA)
         if node.get_weights_by_keys(GAMMA) is None:
             std_output = 1.0
@@ -92,7 +92,7 @@ def _get_mean_std_outputs(node: BaseNode,
             mean_output = 0.0
     else:
         next_node_list = graph.get_next_nodes(node)
-        bn_nodes = [bn_node for bn_node in next_node_list if bn_node.type == BatchNormalization]
+        bn_nodes = [bn_node for bn_node in next_node_list if bn_node.is_match_type(BatchNormalization)]
         if len(bn_nodes) != 0:
             bn_node = bn_nodes[0]
             moving_variance = bn_node.get_weights_by_keys(MOVING_VARIANCE)

model_compression_toolkit/core/keras/pruning/pruning_keras_implementation.py

@@ -209,10 +209,9 @@ def _is_keras_node_pruning_section_edge(node: BaseNode) -> bool:
     """
 
     # Check if the node is a Conv2D or Conv2DTranspose layer with groups set to 1.
-    if node.type in [keras.layers.Conv2D, keras.layers.Conv2DTranspose]:
+    if node.is_match_type(keras.layers.Conv2D) or node.is_match_type(keras.layers.Conv2DTranspose):
         return node.framework_attr[GROUPS] == 1
-    return node.type == keras.layers.Dense
-
+    return node.is_match_type(keras.layers.Dense)
 
 
 def _prune_keras_edge_node(node: BaseNode,
@@ -250,9 +249,9 @@ def _prune_keras_edge_node(node: BaseNode,
 
     if not is_exit_node:
         # Update 'filters' or 'units' attributes for entry node Conv2D/Conv2DTranspose layers.
-        if node.type in [keras.layers.Conv2D, keras.layers.Conv2DTranspose]:
+        if node.is_match_type(keras.layers.Conv2D) or node.is_match_type(keras.layers.Conv2DTranspose):
             node.framework_attr[FILTERS] = int(np.sum(mask))
-        elif node.type == keras.layers.Dense:
+        elif node.is_match_type(keras.layers.Dense):
             node.framework_attr[UNITS] = int(np.sum(mask))
 
     if is_exit_node:

model_compression_toolkit/core/keras/reader/common.py

@@ -43,7 +43,7 @@ def is_node_an_input_layer(node: BaseNode) -> bool:
         Whether the node represents an input layer or not.
     """
     if isinstance(node, BaseNode):
-        return node.type == InputLayer
+        return node.is_match_type(InputLayer)
     elif isinstance(node, KerasNode):
         return isinstance(node.layer, InputLayer)
     else:
@@ -60,7 +60,7 @@ def is_node_a_model(node: BaseNode) -> bool:
         Whether the node represents a Keras model or not.
     """
     if isinstance(node, BaseNode):
-        return node.type in [Functional, Sequential]
+        return node.is_match_type(Functional) or node.is_match_type(Sequential)
     elif isinstance(node, KerasNode):
         return isinstance(node.layer, Functional) or isinstance(node.layer, Sequential)
     else:

model_compression_toolkit/core/keras/reader/node_builder.py

@@ -41,7 +41,7 @@ layers = keras.layers
 
 REUSED_IDENTIFIER = '_reused_'
 
-is_const = lambda x: isinstance(x, (tf.Variable, tf.Tensor, np.ndarray))
+is_const = lambda x: isinstance(x, (tf.Variable, tf.Tensor, np.ndarray, float))
 is_tensor = lambda x: isinstance(x, KerasTensor)
 
 
@@ -61,18 +61,36 @@ def get_kwargs2index(tfoplambda_layer: TFOpLambda) -> Dict[str, int]:
     """
     Positional weights are saved according to their index in the node's call arguments, so
     need to know the function arguments' names in case the weights are in the kwargs.
+
+    Note: the kwargs2index dictionary is initialized manually (and not with tf_inspect) so
+    it will only include the arguments that may contain constants. For example, we don't
+    want the transpose_a attribute of tf.matmul to be saved as a constant.
+
+    Every operation we add support to, needs to be added here.
+
     Args:
         tfoplambda_layer: TFOpLambda layer.
 
     Returns:
         A dictionary with argument number and index: {arg_name: arg_index}.
     """
-    if tfoplambda_layer.function in [tf.add, tf.subtract, tf.divide, tf.truediv, tf.multiply, tf.pow,
-                                     tf.matmul, tf.image.crop_and_resize, tf.image.combined_non_max_suppression] or \
-            tfoplambda_layer.symbol in ['__operators__.add', 'math.add', 'math.multiply', 'linalg.matmul', 'concat']:
-        return {arg_name: i for i, arg_name in enumerate(tf_inspect.getfullargspec(tfoplambda_layer.function).args)}
-    else:
-        return {}
+    kwargs2index = {tf.add: {'x': 0, 'y': 1},
+                    tf.subtract: {'x': 0, 'y': 1},
+                    tf.divide: {'x': 0, 'y': 1},
+                    tf.truediv: {'x': 0, 'y': 1},
+                    tf.multiply: {'x': 0, 'y': 1},
+                    tf.pow: {'x': 0, 'y': 1},
+                    tf.matmul: {'a': 0, 'b': 1}}.get(tfoplambda_layer.function)
+    if not kwargs2index:
+        # In TF 2.15 the function attribute is different and doesn't match the original
+        # operation object we use. Therefore, we extract kwargs2index with the symbol.
+        kwargs2index = {'__operators__.add': {'x': 0, 'y': 1},
+                        'math.add': {'x': 0, 'y': 1},
+                        'math.multiply': {'x': 0, 'y': 1},
+                        'linalg.matmul': {'a': 0, 'b': 1},
+                        'concat': {'values': 0}}.get(tfoplambda_layer.symbol, {})
+
+    return kwargs2index
 
 
 def build_node(node: KerasNode,
@@ -154,8 +172,9 @@ def build_node(node: KerasNode,
             if is_const(v) or (keras_layer.function in [tf.add, tf.multiply, tf.subtract, tf.divide, tf.truediv, tf.pow,
                                                         tf.matmul] and
                                isinstance(v, (tuple, list))):
-                weights.update({kwarg2index[k]: to_numpy(v, is_single_tensor=True)})
-                weight_keys.append(k)
+                if k in kwarg2index:
+                    weights.update({kwarg2index[k]: to_numpy(v, is_single_tensor=True)})
+                    weight_keys.append(k)
         # remove weights and KerasTensors and weights from op_call_kwargs
         op_call_kwargs = {k: v for k, v in op_call_kwargs.items()
                           if not (kwarg2index.get(k) in weights or is_tensor(v))}

model_compression_toolkit/core/keras/tf_tensor_numpy.py

@@ -40,7 +40,7 @@ def to_tf_tensor(tensor):
         Logger.critical(f'Unsupported type for conversion to TF tensor: {type(tensor)}.')
 
 
-def tf_tensor_to_numpy(tensor: Union[List, Tuple, np.ndarray, tf.Tensor],
+def tf_tensor_to_numpy(tensor: Union[List, Tuple, np.ndarray, tf.Tensor, float],
                        is_single_tensor=False) -> np.ndarray:
     """
     Convert a TF tensor to a Numpy array.
@@ -65,6 +65,9 @@ def tf_tensor_to_numpy(tensor: Union[List, Tuple, np.ndarray, tf.Tensor],
         else:
             return (tf_tensor_to_numpy(t) for t in tensor)
     elif isinstance(tensor, tf.Tensor):
-        return tensor.numpy()
+        np_tensor = tensor.numpy()
+        return np.array([np_tensor]) if np.isscalar(np_tensor) else np_tensor
+    elif isinstance(tensor, float):
+        return np.array([tensor])
     else:
         Logger.critical(f'Unsupported type for conversion to Numpy array: {type(tensor)}.')

model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -33,26 +33,31 @@ from model_compression_toolkit.core.pytorch.pytorch_device_config import get_wor
 from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlaceHolder
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from mct_quantizers.common.constants import ACTIVATION_HOLDER_QUANTIZER
+from mct_quantizers import PytorchQuantizationWrapper
 
 
 def _build_input_tensors_list(node: BaseNode,
                               graph: Graph,
                               inputs: Tuple[Any],
-                              node_to_output_tensors_dict: Dict[BaseNode, List]) -> List[List]:
+                              node_to_output_tensors_dict: Dict[BaseNode, List],
+                              is_op_quantize_wrapper: bool) -> List[List]:
     """
-    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.
+    Given a node, build a list of input tensors the node gets. The list is built based on the
+    node's incoming edges, previous nodes' output tensors and the node's positional weights.
+    Positional weights aren't used if the node's op is PytorchQuantizationWrapper, since it's
+    positional weights are already in the wrapper.
 
     Args:
         node: Node to build its input tensors list.
         graph: Graph the node is in.
-        inputs: list of input tensors to model
+        inputs: list of input tensors to model.
         node_to_output_tensors_dict: A dictionary from a node to its output tensors.
+        is_op_quantize_wrapper: Whether the func_op is a PytorchQuantizationWrapper or not.
 
     Returns:
         A list of the node's input tensors.
     """
-    if node.type == DummyPlaceHolder:
+    if node.is_match_type(DummyPlaceHolder):
         input_tensors = [inputs[graph.get_inputs().index(node)]]
     else:
         input_tensors = []
@@ -62,7 +67,8 @@ def _build_input_tensors_list(node: BaseNode,
             _input_tensors = node_to_output_tensors_dict[ie.source_node]
             input_tensors.append(_input_tensors)
         input_tensors = [tensor for tensor_list in input_tensors for tensor in tensor_list]  # flat list of lists
-        input_tensors = node.insert_positional_weights_to_input_list(input_tensors)
+        if not is_op_quantize_wrapper:
+            input_tensors = node.insert_positional_weights_to_input_list(input_tensors)
         # convert inputs from positional weights (numpy arrays) to tensors. Must handle each element in the
         # list separately, because in FX the tensors are FX objects and fail to_torch_tensor
         input_tensors = [to_torch_tensor(t) if isinstance(t, np.ndarray) else t
@@ -70,22 +76,27 @@ def _build_input_tensors_list(node: BaseNode,
     return input_tensors
 
 
-def _merge_inputs(_node, input_tensors: List, op_call_args: List) -> List:
+def _merge_inputs(_node: BaseNode, input_tensors: List, op_call_args: List,
+                  is_op_quantize_wrapper: bool) -> List:
     """
-    Merge input tensors list with op_call_args, according to correct order
+    Merge input tensors list with op_call_args, according to correct order.
 
     Args:
-        _node: The node the inputs are for
+        _node: The node the inputs are for.
         input_tensors: activation input tensors to node.
-        op_call_args: framework node call args
+        op_call_args: framework node call args.
+        is_op_quantize_wrapper: Whether the func_op is a PytorchQuantizationWrapper or not.
     Returns:
-        Combined list of input_tensors and op_call_args
+        Combined list of input_tensors and op_call_args.
     """
     if isinstance(_node, FunctionalNode) and _node.tensor_input_indices:
-        assert len(_node.tensor_input_indices) == len(input_tensors), 'Mismatch between input tensors and indices'
         _input_list = op_call_args.copy()
-        for i, t in zip(_node.tensor_input_indices, input_tensors):
-            _input_list.insert(i, t)
+        if is_op_quantize_wrapper:
+            _input_list = input_tensors + _input_list
+        else:
+            assert len(_node.tensor_input_indices) == len(input_tensors), 'Mismatch between input tensors and indices'
+            for i, t in zip(_node.tensor_input_indices, input_tensors):
+                _input_list.insert(i, t)
     else:
         _input_list = input_tensors + op_call_args
 
@@ -118,7 +129,8 @@ def _run_operation(n: BaseNode,
     if isinstance(n, FunctionalNode) and n.inputs_as_list:
         out_tensors_of_n_float = op_func(input_tensors, *op_call_args, **functional_kwargs)
     else:
-        out_tensors_of_n_float = op_func(*_merge_inputs(n, input_tensors, op_call_args), **functional_kwargs)
+        merged_inputs = _merge_inputs(n, input_tensors, op_call_args, isinstance(op_func, PytorchQuantizationWrapper))
+        out_tensors_of_n_float = op_func(*merged_inputs, **functional_kwargs)
 
     # Add a fake quant node if the node has an activation threshold.
     out_tensors_of_n = out_tensors_of_n_float
@@ -279,12 +291,12 @@ class PytorchModel(torch.nn.Module):
         node_to_output_tensors_dict_float = dict()
         configurable_nodes = self.graph.get_configurable_sorted_nodes_names(DEFAULT_PYTORCH_INFO)
         for node in self.node_sort:
+            op_func = self._get_op_func(node, configurable_nodes)
             input_tensors = _build_input_tensors_list(node,
                                                       self.graph,
                                                       args,
-                                                      node_to_output_tensors_dict)
-
-            op_func = self._get_op_func(node, configurable_nodes)
+                                                      node_to_output_tensors_dict,
+                                                      isinstance(op_func, PytorchQuantizationWrapper))
             use_activation_quantization, activation_quantization_fn = self._get_activation_quantization_fn(node)
 
             # Run node operation and fetch outputs
@@ -326,15 +338,16 @@ class PytorchModel(torch.nn.Module):
         """
         return getattr(self, node.name)
 
-    def _get_activation_quantization_fn(self, node) -> Tuple[bool, bool, Callable]:
+    def _get_activation_quantization_fn(self, node) -> Tuple[bool, Callable]:
         """
         Get activation quantization parameters for this node.
 
         Args:
             node: Node from which to extract the activation quantization parameters.
 
-        Returns: Flag to indicate if we quantize activations, flag to indicate if we quantize activations
-        using a quantization holder and a quantization function to use for the node's activations.
+        Returns:
+            Flag to indicate if we quantize activations using a quantization holder and a quantization
+            function to use for the node's activations.
         """
         activation_quantization_holder = self.node_to_activation_quantization_holder.get(node.name)
         use_activation_quantization = node.is_activation_quantization_enabled()

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/batchnorm_folding.py

@@ -62,11 +62,11 @@ def update_kernel_for_bn_folding_fn(conv_node: BaseNode,
     Returns:
         The modified convolution node's weight/kernel/
     """
-    if conv_node.type == ConvTranspose2d:
+    if conv_node.is_match_type(ConvTranspose2d):
         _scale = weights_scale[None, :, None, None]
     else:
         _scale = weights_scale[:, None, None, None]
-    if conv_node.type == ConvTranspose2d and conv_node.framework_attr[GROUPS] > 1:
+    if conv_node.is_match_type(ConvTranspose2d) and conv_node.framework_attr[GROUPS] > 1:
         # PyTorch ConvTranspose2d kernel with groups stacks groups on in_channels axis, so need to reshape the kernel
         # so the groups are stacked on the out_channels axis to match the scale vector (then reshape back to original
         # shape)
@@ -93,10 +93,10 @@ def update_weights_for_bn_forward_folding_fn(conv_node: BaseNode,
     Returns:
         The modified convolution node's weight/kernel/
     """
-    if conv_node.type == Conv2d and conv_node.framework_attr[GROUPS] > 1:
+    if conv_node.is_match_type(Conv2d) and conv_node.framework_attr[GROUPS] > 1:
         bias_update = (kernel * bias_factor[:, None, None, None]).flatten()
         _scale = weights_scale[:, None, None, None]
-    elif conv_node.type == ConvTranspose2d:
+    elif conv_node.is_match_type(ConvTranspose2d):
         bias_update = (kernel * bias_factor[:, None, None, None]).sum(axis=0).flatten()
         _scale = weights_scale[:, None, None, None]
     else:
@@ -125,8 +125,8 @@ def is_group_conv_fn(node: BaseNode) -> bool:
     Returns:
         True if the node is a group convolution, else False
     """
-    return node.type in [Conv2d, ConvTranspose2d] and \
-           node.framework_attr[GROUPS] not in [node.framework_attr[IN_CHANNELS], 1]
+    return (node.is_match_type(Conv2d) or node.is_match_type(ConvTranspose2d)) and \
+        node.framework_attr[GROUPS] not in [node.framework_attr[IN_CHANNELS], 1]
 
 
 def get_foldable_node_type_and_validity_fn(node: BaseNode) -> [bool, bool]:
@@ -140,8 +140,8 @@ def get_foldable_node_type_and_validity_fn(node: BaseNode) -> [bool, bool]:
         is_bn: True if the node is a batch norm, else False
         is_dw_valid: True if the node is a dw-convolution valid for folding or a batch-norm node, else False
     """
-    is_bn = node.type is BatchNorm2d
-    is_dw = node.type is Conv2d and node.framework_attr[GROUPS] == node.framework_attr[IN_CHANNELS]
+    is_bn = node.is_match_type(BatchNorm2d)
+    is_dw = node.is_match_type(Conv2d) and node.framework_attr[GROUPS] == node.framework_attr[IN_CHANNELS]
     is_dw_valid = is_dw and np.all(np.array(node.get_weights_by_keys(KERNEL).shape[2:]) == 1)
     return is_bn, is_dw_valid
 

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/const_holder_conv.py

@@ -48,9 +48,9 @@ class FunctionalConvSubstitution(common.BaseSubstitution):
             Graph after applying the substitution.
         """
         # Set new layer
-        if func_node.type == conv2d:
+        if func_node.is_match_type(conv2d):
             new_layer = Conv2d
-        elif func_node.type == conv_transpose2d:
+        elif func_node.is_match_type(conv_transpose2d):
             new_layer = ConvTranspose2d
         else:
             Logger.critical(f'Substitution filter mismatch. Layer {func_node.type}. Must be {type(Conv2d)} or {type(ConvTranspose2d)}.')  # pragma: no cover

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

@@ -53,7 +53,7 @@ def conv2d_collapsing_fn(first_node: BaseNode,
     Returns:
         The modified layer node's weights: kernel, bias
     """
-    if first_node.type == Conv2d and second_node.type == Conv2d:
+    if first_node.is_match_type(Conv2d) and second_node.is_match_type(Conv2d):
         # Get nodes attributes
         kernel1 = first_node.get_weights_by_keys(kernel_str)
         kernel2 = second_node.get_weights_by_keys(kernel_str)

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

@@ -76,17 +76,17 @@ class ReLUBoundToPowerOfTwo(common.BaseSubstitution):
         second_op2d_node = nodes_list[2]
 
         # only act on bound relu with not POT max value and 0 min value
-        if non_linear_node.type == ReLU6:
+        if non_linear_node.is_match_type(ReLU6):
             scale_factor = 6.0 / self.threshold
             non_linear_node.layer_class = Hardtanh
             non_linear_node.framework_attr[INPLACE] = False
             non_linear_node.framework_attr[HARDTANH_MIN_VAL] = 0.0
             non_linear_node.framework_attr[HARDTANH_MAX_VAL] = self.threshold
-        elif non_linear_node.type == relu6:
+        elif non_linear_node.is_match_type(relu6):
             scale_factor = 6.0 / self.threshold
             non_linear_node.functional_op = hardtanh
             non_linear_node.functional_op.__defaults__ = (0.0, self.threshold, False)
-        elif non_linear_node.type == Hardtanh:
+        elif non_linear_node.is_match_type(Hardtanh):
             if (non_linear_node.framework_attr[HARDTANH_MIN_VAL] == 0.0) and not \
                     (np.log2(non_linear_node.framework_attr[HARDTANH_MAX_VAL]).astype(int) -
                      np.log2(non_linear_node.framework_attr[HARDTANH_MAX_VAL]) == 0):
@@ -94,7 +94,7 @@ class ReLUBoundToPowerOfTwo(common.BaseSubstitution):
                 non_linear_node.framework_attr[HARDTANH_MAX_VAL] = self.threshold
             else:
                 return graph
-        elif non_linear_node.type == hardtanh:
+        elif non_linear_node.is_match_type(hardtanh):
             if (non_linear_node.framework_attr[HARDTANH_MIN_VAL] == 0.0) and not \
                     (np.log2(non_linear_node.framework_attr[HARDTANH_MAX_VAL]).astype(int) -
                      np.log2(non_linear_node.framework_attr[HARDTANH_MAX_VAL]) == 0):

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/residual_collapsing.py

@@ -46,7 +46,7 @@ def residual_collapsing_fn(first_node: BaseNode,
     Returns:
         The modified layer node's weights: kernel
     """
-    if first_node.type == Conv2d:
+    if first_node.is_match_type(Conv2d):
         # Get nodes attributes
         kernel = first_node.get_weights_by_keys(kernel_str)
         (Cout, Cin, kH, kW) = kernel.shape

model_compression_toolkit/core/pytorch/pruning/pruning_pytorch_implementation.py

@@ -76,9 +76,9 @@ class PruningPytorchImplementation(PytorchImplementation, PruningFrameworkImplem
         pruned_parameters = {}
         mask_bool = output_mask.astype(bool)
         node.weights = pruned_parameters
-        if node.type == torch.nn.BatchNorm2d:
+        if node.is_match_type(torch.nn.BatchNorm2d):
             node.framework_attr[NUM_FEATURES] = int(np.sum(input_mask))
-        elif node.type == torch.nn.PReLU:
+        elif node.is_match_type(torch.nn.PReLU):
             if node.framework_attr[NUM_PARAMETERS] > 1:
                 node.framework_attr[NUM_PARAMETERS] = int(np.sum(input_mask))
             else:
@@ -227,9 +227,9 @@ def _is_pytorch_node_pruning_section_edge(node: BaseNode) -> bool:
     """
 
     # Check if the node is a Conv2D or Conv2DTranspose layer with groups set to 1.
-    if node.type in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:
+    if node.is_match_type(torch.nn.Conv2d) or node.is_match_type(torch.nn.ConvTranspose2d):
         return node.framework_attr[GROUPS] == 1
-    return node.type == torch.nn.Linear
+    return node.is_match_type(torch.nn.Linear)
 
 
 def _prune_pytorch_edge_node(node: BaseNode,
@@ -268,18 +268,18 @@ def _prune_pytorch_edge_node(node: BaseNode,
     if not is_exit_node:
         # Update 'out_channels' or 'out_features' attributes for entry nodes
         # Conv2d,ConvTranspose2d / Linear layers.
-        if node.type in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:
+        if node.is_match_type(torch.nn.Conv2d) or node.is_match_type(torch.nn.ConvTranspose2d):
             node.framework_attr[OUT_CHANNELS] = int(np.sum(mask))
-        elif node.type == torch.nn.Linear:
+        elif node.is_match_type(torch.nn.Linear):
             node.framework_attr[OUT_FEATURES] = int(np.sum(mask))
         else:
             Logger.critical(f"{node.type} is currently not supported"
                              f"as an edge node in a pruning section")
 
     if is_exit_node:
-        if node.type in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:
+        if node.is_match_type(torch.nn.Conv2d) or node.is_match_type(torch.nn.ConvTranspose2d):
             node.framework_attr[IN_CHANNELS] = int(np.sum(mask))
-        elif node.type == torch.nn.Linear:
+        elif node.is_match_type(torch.nn.Linear):
             node.framework_attr[IN_FEATURES] = int(np.sum(mask))
         else:
             Logger.critical(f"{node.type} is currently not supported"

model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -398,8 +398,8 @@ class PytorchImplementation(FrameworkImplementation):
         Returns: True if the node should be considered an interest point, False otherwise.
         """
 
-        if node.type in [Conv2d, Linear, ConvTranspose2d, Sigmoid, sigmoid, Softmax, softmax, operator.add, add, cat,
-                         operator.concat]:
+        if any([node.is_match_type(_type) for _type in [Conv2d, Linear, ConvTranspose2d, Sigmoid, sigmoid, Softmax,
+                                                        softmax, operator.add, add, cat, operator.concat]]):
             return True
         return False
 
@@ -464,12 +464,12 @@ class PytorchImplementation(FrameworkImplementation):
         kernel_shape = node.get_weights_by_keys(fw_info.get_kernel_op_attributes(node.type)[0]).shape
         output_channel_axis, input_channel_axis = fw_info.kernel_channels_mapping.get(node.type)
 
-        if node.type is Conv2d or node.type is ConvTranspose2d:
+        if node.is_match_type(Conv2d) or node.is_match_type(ConvTranspose2d):
             # (C_out * W_out * H_out) * C_in * (W_kernel * H_kernel)
             return np.prod([x for x in output_shape if x is not None]) * \
                    kernel_shape[input_channel_axis] * \
                    (kernel_shape[0] * kernel_shape[1])
-        elif node.type is Linear:
+        elif node.is_match_type(Linear):
             # IN * OUT
             return kernel_shape[0] * kernel_shape[1]
         else:
@@ -552,7 +552,6 @@ class PytorchImplementation(FrameworkImplementation):
         Returns:
             weight_quantizers: A dictionary between a weight's name to its quantizer.
             activation_quantizers: A list of activations quantization, one for each layer output.
-
         """
 
         return get_inferable_quantizers(node,

model_compression_toolkit/core/pytorch/pytorch_node_prior_info.py

@@ -62,7 +62,7 @@ def _get_mean_std_outputs(node: BaseNode,
     """
     mean_output, std_output = None, None
 
-    if node.type == BatchNorm2d:
+    if node.is_match_type(BatchNorm2d):
         mean_output = node.get_weights_by_keys(BETA)
         if node.get_weights_by_keys(GAMMA) is None:
             std_output = 1.0
@@ -72,7 +72,7 @@ def _get_mean_std_outputs(node: BaseNode,
             mean_output = 0.0
     else:
         next_node_list = graph.get_next_nodes(node)
-        bn_nodes = [bn_node for bn_node in next_node_list if bn_node.type == BatchNorm2d]
+        bn_nodes = [bn_node for bn_node in next_node_list if bn_node.is_match_type(BatchNorm2d)]
         if len(bn_nodes) != 0:
             bn_node = bn_nodes[0]
             moving_variance = bn_node.get_weights_by_keys(MOVING_VARIANCE)

model_compression_toolkit/exporter/model_wrapper/keras/builder/fully_quantized_model_builder.py

@@ -42,8 +42,12 @@ if FOUND_TF:
         """
         weights_quantizers, _ = fw_impl.get_inferable_quantizers(node)
         if len(weights_quantizers) > 0:
+            # for positional weights we need to extract the weight's value.
+            weights_values = {attr: node.get_weights_by_keys(attr)
+                              for attr in weights_quantizers if isinstance(attr, int)}
             return KerasQuantizationWrapper(layer,
-                                            weights_quantizers)
+                                            weights_quantizers,
+                                            weights_values)
         return layer
 
 

model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py

@@ -29,7 +29,7 @@ if FOUND_TORCH:
 
     def fully_quantized_wrapper(node: common.BaseNode,
                                 module: torch.nn.Module,
-                                fw_impl) -> Union[torch.nn.Module,PytorchQuantizationWrapper]:
+                                fw_impl) -> Union[torch.nn.Module, PytorchQuantizationWrapper]:
         """
         A function which takes a computational graph node and a pytorch module and
         perform the quantization wrapping
@@ -37,20 +37,26 @@ if FOUND_TORCH:
         Args:
             node: A node of mct graph.
             module: A Pytorch module
+            fw_impl: FrameworkImplementation object with a specific framework methods implementation.
         Returns: Wrapped layer
 
         """
         weight_quantizers, _ = fw_impl.get_inferable_quantizers(node)
         if len(weight_quantizers) > 0:
-            return PytorchQuantizationWrapper(module, weight_quantizers)
+            # for positional weights we need to extract the weight's value.
+            weights_values = {attr: fw_impl.to_tensor(node.get_weights_by_keys(attr))
+                              for attr in weight_quantizers if isinstance(attr, int)}
+            return PytorchQuantizationWrapper(module, weight_quantizers, weights_values)
         return module
 
+
     def get_activation_quantizer_holder(node: BaseNode, fw_impl) -> Callable:
         """
         Retrieve a PytorchActivationQuantizationHolder layer to use for activation quantization of a node.
         If the layer is not supposed to be wrapped with an activation quantizer - return None.
         Args:
             node: Node to attach a PytorchActivationQuantizationHolder to its output.
+            fw_impl: FrameworkImplementation object with a specific framework methods implementation.
         Returns:
             A PytorchActivationQuantizationHolder module for the node's activation quantization.
         """
@@ -64,6 +70,7 @@ if FOUND_TORCH:
             f'PytorchActivationQuantizationHolder supports a single quantizer but {len(activation_quantizers)} quantizers '
             f'were found for node {node}')
 
+
     def get_exportable_pytorch_model(graph: Graph):
         """
         Convert graph to fully quantized PyTorch model.