mct-nightly 2.1.0.20240622.419__py3-none-any.whl → 2.1.0.20240624.520__py3-none-any.whl

{mct_nightly-2.1.0.20240622.419.dist-info → mct_nightly-2.1.0.20240624.520.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.1.0.20240622.419
+Version: 2.1.0.20240624.520
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct_nightly-2.1.0.20240622.419.dist-info → mct_nightly-2.1.0.20240624.520.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=bxCfxpPll_b3K-MWuleZf0LpbrzZKqMQFrBpMlfRQDk,1573
+model_compression_toolkit/__init__.py,sha256=9NcQg8E0BkwMF32oeqvh_b8tuvTkx4OAmKJ_1q74DmE,1573
 model_compression_toolkit/constants.py,sha256=9pVleMwnhlM4QwIL2HcEq42I1uF4rlSw63RUjkxOF4w,3923
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -7,7 +7,7 @@ model_compression_toolkit/core/__init__.py,sha256=TrRgkWpT1AN2Faw1M_1HXyJkJnbxfn
 model_compression_toolkit/core/analyzer.py,sha256=X-2ZpkH1xdXnISnw1yJvXnvV-ssoUh-9LkLISSWNqiY,3691
 model_compression_toolkit/core/graph_prep_runner.py,sha256=kM70wmNG3yMFiGQc0uO0wn9j4ZbSWxUEykpxDK55doc,10567
 model_compression_toolkit/core/quantization_prep_runner.py,sha256=0ga95vh_ZXO79r8FB26L5GIZKHkG98wq1hMsNH1bIeU,6453
-model_compression_toolkit/core/runner.py,sha256=kMr8pK2Z_F1Fl3uHf6ymeNKEH1NaPWQjEGEqM7sRn04,12654
+model_compression_toolkit/core/runner.py,sha256=4TtOgyNb4cXr52dOlDqYxLm3rnLR6uHPDNoZiEFL9XA,12655
 model_compression_toolkit/core/common/__init__.py,sha256=Wh127PbXcETZX_d1PQqZ71ETK3J9XO5A-HpadGUbj6o,1447
 model_compression_toolkit/core/common/base_substitutions.py,sha256=xDFSmVVs_iFSZfajytI0cuQaNRNcwHX3uqOoHgVUvxQ,1666
 model_compression_toolkit/core/common/framework_implementation.py,sha256=8b6M1GcUR9bDgoxwqyNP8C6KSU9OTQ5hIk20Y74eLPo,20896
@@ -33,7 +33,7 @@ model_compression_toolkit/core/common/graph/__init__.py,sha256=Xr-Lt_qXMdrCnnOaU
 model_compression_toolkit/core/common/graph/base_graph.py,sha256=lmIw0srKiwCvz7KWqfwKTxyQHDy3s6rWMIXzFAa1UMo,38326
 model_compression_toolkit/core/common/graph/base_node.py,sha256=X_0zqHrKYAsmnj9tAKjVYasbFcZD8OHpjdiMj9ugQs0,29436
 model_compression_toolkit/core/common/graph/edge.py,sha256=buoSEUZwilWBK3WeBKpJ-GeDaUA1SDdOHxDpxU_bGpk,3784
-model_compression_toolkit/core/common/graph/functional_node.py,sha256=_6HsBeLlrpLvXhLPRJswcyDa4z16-O3xzHzGuv46zBc,3897
+model_compression_toolkit/core/common/graph/functional_node.py,sha256=BbxQ-WRk4R-5hbpQDBANkhRRTkaG7eogeiJwLfLb_EU,3950
 model_compression_toolkit/core/common/graph/graph_matchers.py,sha256=CrDoHYq4iPaflgJWmoJ1K4ziLrRogJvFTVWg8P0UcDU,4744
 model_compression_toolkit/core/common/graph/graph_searches.py,sha256=2oKuW6L8hP-oL0lFO9PhQFt9fEFgVJwpc1u4fHExAtE,5128
 model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py,sha256=3el-A7j1oyoo1_9zq3faQp7IeRsFXFCvnrb3zZFXpU0,9803
@@ -222,7 +222,7 @@ model_compression_toolkit/core/pytorch/back2framework/factory_model_builder.py,s
 model_compression_toolkit/core/pytorch/back2framework/float_model_builder.py,sha256=tLrlUyYhxVKVjkad1ZAtbRra0HedB3iVfIkZ_dYnQ-4,3419
 model_compression_toolkit/core/pytorch/back2framework/instance_builder.py,sha256=BBHBfTqeWm7L3iDyPBpk0jxvj-rBg1QWI23imkjfIl0,1467
 model_compression_toolkit/core/pytorch/back2framework/mixed_precision_model_builder.py,sha256=D7lU1r9Uq_7fdNuKk2BMF8ho5GrsY-8gyGN6yYoHaVg,15060
-model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py,sha256=Pvpkirt3ziWEXDEspgOhR8ALf-XAZUh-78IkXg9YMWs,18830
+model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py,sha256=b3RJ9XpbN2XXlCXEVjxLg3NenmtFfnp_UBRKDIEka8A,18698
 model_compression_toolkit/core/pytorch/back2framework/quantized_model_builder.py,sha256=qZNNOlNTTV4ZKPG3q5GDXkIVTPUEr8dvxAS_YiMORmg,3456
 model_compression_toolkit/core/pytorch/back2framework/quantization_wrapper/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/core/pytorch/back2framework/quantization_wrapper/quantized_layer_wrapper.py,sha256=q2JDw10NKng50ee2i9faGzWZ-IydnR2aOMGSn9RoZmc,5773
@@ -233,9 +233,9 @@ model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/batchno
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/batchnorm_reconstruction.py,sha256=B7aC2TZNrQJ2oQVGBFhKAVqdUU5lYVJSMmwKhjxOHWk,2822
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/batchnorm_refusing.py,sha256=JDWOaNwYrZG0zTwd3HwoZUM3tKu7zPbzLOrqNQsu8xA,2162
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/concat_threshold_update.py,sha256=SBrR24ZAnWPftLinv4FuIqdBGjfYtfXbYQJN5mgy5V4,2861
-model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/const_holder_conv.py,sha256=Fs2YQBD4KJV-pGLOMqm-p485bfq2JDYgCzFroRljCoM,3933
-model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_batch_norm.py,sha256=iX8bLHtw2osP42-peNLTRmbpX3cUxdGsAbEfw7NLpx0,3935
-model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_layer_norm.py,sha256=zKSgtVw_P9fUvdq4e7P9yaLDPG_vZ0cecM9sVPtm1ns,3799
+model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/const_holder_conv.py,sha256=iTuP1hjuTZTGcE7izfs_UOWBGeEBFRvRIU4QCh-b21M,4627
+model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_batch_norm.py,sha256=7GZY7lU3LUUaO5iiccHkUP62PB0QeGAGOZdUSGMkFBY,4450
+model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_layer_norm.py,sha256=XhiLVcnCc_gF-6mjxbf9C4bYg5YL_GCvDJmcdLkBNAg,4151
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/linear_collapsing.py,sha256=CXSMASpc_Zed3BJ2CsER69zKxE6ncFvvKQWDO1JxKYI,5849
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/multi_head_attention_decomposition.py,sha256=VNg-VgzCxSyqy2J3neEPl6U0SPO8UIVU_T47bGhz4FE,38459
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/relu_bound_to_power_of_2.py,sha256=q1a3HieQtaOmWG2WGXp6GHYAvxa3CZ9dJUx9dqMAsS8,5695
@@ -261,7 +261,7 @@ model_compression_toolkit/core/pytorch/quantizer/__init__.py,sha256=Rf1RcYmelmdZ
 model_compression_toolkit/core/pytorch/quantizer/fake_quant_builder.py,sha256=D8_CEuFqKAhbUgKaRw7Jlxo0zlqgPTMu6CIIIM4LfS0,7045
 model_compression_toolkit/core/pytorch/quantizer/lut_fake_quant.py,sha256=uyeBtNokyDUikk-YkDP_mN_2DX0J5oPm3kSfdSUT2Ck,4420
 model_compression_toolkit/core/pytorch/reader/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
-model_compression_toolkit/core/pytorch/reader/graph_builders.py,sha256=LiGV-ZqlhxN1evpM-ur2dDVPowhrLwO7JZa7AGPftSk,12913
+model_compression_toolkit/core/pytorch/reader/graph_builders.py,sha256=ESL8k7RLZogTyG_oTTFDmm4RauZvx2gU-UvnOnEsH6Q,15948
 model_compression_toolkit/core/pytorch/reader/node_holders.py,sha256=TaolORuwBZEddWe-q0Mg79Nmswz-Sq3-9-4o8UxFQ50,1028
 model_compression_toolkit/core/pytorch/reader/reader.py,sha256=GEJE0QX8XJFWbYCkbRBtzttZtmmuoACLx8gw9KyAQCE,6015
 model_compression_toolkit/core/pytorch/statistics_correction/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
@@ -425,7 +425,7 @@ model_compression_toolkit/target_platform_capabilities/target_platform/targetpla
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/attribute_filter.py,sha256=jfhszvuD2Fyy6W2KjlLzXBQKFzTqGAaDZeFVr4-ONQw,8776
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/current_tpc.py,sha256=fIheShGOnxWYKqT8saHpBJqOU5RG_1Hp9qHry7IviIw,2115
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/layer_filter_params.py,sha256=Cl6-mACpje2jM8RJkibbqE3hvTkFR3r26-lW021mIiA,4019
-model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/operations_to_layers.py,sha256=1JN3yvNiJyDfva0tLTH3ej_qORzrQcPz32bSMKl49_0,6720
+model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/operations_to_layers.py,sha256=iZDgHd0SVbgNTT-jtSP0SWsaRGfAJM_p-wpBlBkpRAQ,6723
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/target_platform_capabilities.py,sha256=KP8IWlHzkXzVjqIiRtAW6sTYyHJ2wVFFX4hMt_N6o3s,9910
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/target_platform_capabilities_component.py,sha256=FvrYI0Qy7DCmDp2gyUYyCZq5pY84JgLtJqSIiVTJ8Ss,1030
 model_compression_toolkit/target_platform_capabilities/tpc_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -517,8 +517,8 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=yrZNVRm2IRU7r7R-hjS2lOQ6wvEEvbeunvf2jKoWjXk,3277
 model_compression_toolkit/xquant/pytorch/similarity_functions.py,sha256=CERxq5K8rqaiE-DlwhZBTUd9x69dtYJlkHOPLB54vm8,2354
 model_compression_toolkit/xquant/pytorch/tensorboard_utils.py,sha256=eyMoXt5o5EnMr6d-rpCwQdX5mAiYiymvbgKv4tf7-a0,4576
-mct_nightly-2.1.0.20240622.419.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.1.0.20240622.419.dist-info/METADATA,sha256=k_tXC3j1eR9oQnfMpjduwljZv8cN1y9RhPwG6P09DqM,19726
-mct_nightly-2.1.0.20240622.419.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-mct_nightly-2.1.0.20240622.419.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.1.0.20240622.419.dist-info/RECORD,,
+mct_nightly-2.1.0.20240624.520.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.1.0.20240624.520.dist-info/METADATA,sha256=0Lh6S3Ea0DK-D1dmGnRH-IwyzULmoho7PC7LXgUL5x0,19726
+mct_nightly-2.1.0.20240624.520.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+mct_nightly-2.1.0.20240624.520.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.1.0.20240624.520.dist-info/RECORD,,

model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.1.0.20240622.000419"
+__version__ = "2.1.0.20240624.000520"

model_compression_toolkit/core/common/graph/functional_node.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, Tuple, Type
+from typing import Dict, Any, Tuple, Type, List, Union
 
 from model_compression_toolkit.constants import FOUND_TF
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
@@ -25,7 +25,7 @@ class FunctionalNode(BaseNode):
                  functional_op: Any = None,
                  inputs_as_list: bool = False,
                  has_activation: bool = True,
-                 tensor_input_allocs = None):
+                 tensor_input_allocs: List[Union[int, str]] = None):
         """
         Init a FunctionalNode object.
 
@@ -44,8 +44,7 @@ class FunctionalNode(BaseNode):
             functional_op: The op the node implements.
             inputs_as_list: Whether to pass the node its input tensors as a list or not when calling the layer.
             has_activation: Whether the node has activations that we might want to quantize.
-            tensor_input_allocs: A list of indices for activation tensors in the node's input tensor list
-
+            tensor_input_allocs: A list of indices and strings for allocations input tensors in the node's args and kwargs.
         """
 
         super().__init__(name,

model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -106,7 +106,7 @@ def _run_operation(n: BaseNode,
                    input_tensors: List,
                    op_func: Any,
                    quantize_node_activation_fn,
-                   use_activation_quantization: bool) -> Tuple[Union[List, torch.Tensor], Union[List, torch.Tensor]]:
+                   use_activation_quantization: bool) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
     """
     Applying the layer (op_func) to the input tensors (input_tensors).
     If quantized is set to True, and the layer's corresponding node (n) has quantization
@@ -126,17 +126,17 @@ def _run_operation(n: BaseNode,
     op_call_args = n.op_call_args if isinstance(n, FunctionalNode) else []
     functional_kwargs = n.op_call_kwargs if isinstance(n, FunctionalNode) else {}
 
-    if not (isinstance(n, FunctionalNode) and isinstance(op_func, PytorchQuantizationWrapper)):
-        # Insert positional weights only when not a quantized functional node, because quantized functional nodes
-        # insert the quantized weights in the wrapper.
+    # Insert positional weights only when not a quantized functional node, because quantized functional nodes
+    # insert the quantized weights in the wrapper.
+    if isinstance(n, FunctionalNode) and isinstance(op_func, PytorchQuantizationWrapper):
+        _tensor_input_allocs = [i for i in n.tensor_input_allocs if i not in n.weights]
+    else:
         input_tensors = n.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, numpy_type=t.dtype) if isinstance(t, np.ndarray) else t
                          for t in input_tensors]
         _tensor_input_allocs = None
-    else:
-        _tensor_input_allocs = [i for i in n.tensor_input_allocs if i not in n.weights]
 
     if isinstance(n, FunctionalNode) and n.inputs_as_list:
         out_tensors_of_n_float = op_func(input_tensors, *op_call_args, **functional_kwargs)
@@ -152,6 +152,8 @@ def _run_operation(n: BaseNode,
             out_tensors_of_n_float = torch.cat(out_tensors_of_n_float, dim=0)
         out_tensors_of_n = quantize_node_activation_fn(out_tensors_of_n_float)
 
+    if not isinstance(out_tensors_of_n, list):
+        out_tensors_of_n, out_tensors_of_n_float = [out_tensors_of_n], [out_tensors_of_n_float]
     return out_tensors_of_n, out_tensors_of_n_float
 
 
@@ -318,12 +320,8 @@ class PytorchModel(torch.nn.Module):
                                                                       quantize_node_activation_fn=activation_quantization_fn,
                                                                       use_activation_quantization=use_activation_quantization)
 
-            if isinstance(out_tensors_of_n, list):
-                node_to_output_tensors_dict.update({node: out_tensors_of_n})
-                node_to_output_tensors_dict_float.update({node: out_tensors_of_n_float})
-            else:
-                node_to_output_tensors_dict.update({node: [out_tensors_of_n]})
-                node_to_output_tensors_dict_float.update({node: [out_tensors_of_n_float]})
+            node_to_output_tensors_dict.update({node: out_tensors_of_n})
+            node_to_output_tensors_dict_float.update({node: out_tensors_of_n_float})
 
         if self.append2output:
             outputs = _generate_outputs(self.append2output,

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

@@ -19,6 +19,7 @@ from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
+from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.pytorch.constants import IN_CHANNELS, OUT_CHANNELS, KERNEL_SIZE, KERNEL, BIAS
 from model_compression_toolkit.core.common import FrameworkInfo
 
@@ -37,7 +38,7 @@ class FunctionalConvSubstitution(common.BaseSubstitution):
 
     def substitute(self,
                    graph: Graph,
-                   func_node: BaseNode) -> Graph:
+                   func_node: FunctionalNode) -> Graph:
         """
         Substitute functional and conv/linear layer with torch layer
         Args:
@@ -60,9 +61,15 @@ class FunctionalConvSubstitution(common.BaseSubstitution):
         # Create new node of layer convolution
         if 1 not in func_node.weights:
             Logger.critical(f'Weight input missing for node {func_node.name}.')  # pragma: no cover
-        weight = func_node.weights[1]
-        bias = func_node.weights.get(2)
-        framework_attr = func_node.framework_attr
+        # Extract index of kernel and bias according to tensor_input_allocs if they were input as kwargs. If
+        # they were input as args, use their fixed positions.
+        weight_index = func_node.tensor_input_allocs.index(KERNEL) if KERNEL in func_node.tensor_input_allocs else 1
+        bias_index = func_node.tensor_input_allocs.index(BIAS) if BIAS in func_node.tensor_input_allocs else 2
+        if weight_index not in func_node.weights:
+            Logger.critical(f'Mismatch between tensor_input_allocs and weight index in node {func_node.name}.')  # pragma: no cover
+        weight = func_node.weights[weight_index]
+        bias = func_node.weights.get(bias_index)
+        framework_attr = func_node.op_call_kwargs
         framework_attr.update({OUT_CHANNELS: weight.shape[out_channel_index]})
         framework_attr.update({IN_CHANNELS: weight.shape[in_channel_index]})
         framework_attr.update({KERNEL_SIZE: weight.shape[2:]})

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

@@ -20,6 +20,7 @@ import torch.nn.functional as F
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode, Graph
+from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.pytorch.constants import *
 from model_compression_toolkit.logger import Logger
 
@@ -37,9 +38,12 @@ class FunctionalBatchNorm(common.BaseSubstitution):
         super().__init__(matcher_instance=bn_node)
 
     @staticmethod
-    def get_attributes_from_weights(node: BaseNode) -> Dict:
+    def get_attributes_from_weights(node: FunctionalNode) -> Dict:
         """
-        convert functional batch_norm positional weights to BatchNorm2d weights
+        Convert functional batch_norm positional weights to BatchNorm2d weights. Extract indices of gamma
+        and beta according to tensor_input_allocs if they were input as kwargs. If they were input as args,
+        use their fixed positions.
+
         Args:
             node: functional batch_norm node.
 
@@ -53,23 +57,22 @@ class FunctionalBatchNorm(common.BaseSubstitution):
                         GAMMA: np.ones(node.weights[1].shape),
                         BETA: np.zeros(node.weights[1].shape)}
 
-        has_weight = WEIGHT not in node.framework_attr
-        has_bias = BIAS not in node.framework_attr
+        # Check if weight and/or bias were not given.
+        if KERNEL in node.tensor_input_allocs:
+            weights_dict[GAMMA] = node.weights[node.tensor_input_allocs.index(KERNEL)]
+        elif KERNEL not in node.op_call_kwargs:
+            weights_dict[GAMMA] = node.weights[3]
 
-        if 3 in node.weights:
-            if has_weight:
-                weights_dict[GAMMA] = node.weights[3]
-            else:
-                weights_dict[BETA] = node.weights[3]
-        if 4 in node.weights:
-            assert has_bias
+        if BIAS in node.tensor_input_allocs:
+            weights_dict[BETA] = node.weights[node.tensor_input_allocs.index(BIAS)]
+        elif BIAS not in node.op_call_kwargs:
             weights_dict[BETA] = node.weights[4]
 
         return weights_dict
 
     def substitute(self,
                    graph: Graph,
-                   node: BaseNode) -> Graph:
+                   node: FunctionalNode) -> Graph:
         """
         Substitute functional.batch_norm and its inputs with BatchNorm2d.
         Args:
@@ -87,10 +90,13 @@ class FunctionalBatchNorm(common.BaseSubstitution):
         bn_node_weights = self.get_attributes_from_weights(node)
         if not bn_node_weights:
             return graph
+        framework_attr = {NUM_FEATURES: out_channels}
+        if EPSILON in node.op_call_kwargs:
+            framework_attr.update({EPSILON: node.op_call_kwargs[EPSILON]})
+        if MOMENTUM in node.op_call_kwargs:
+            framework_attr.update({MOMENTUM: node.op_call_kwargs[MOMENTUM]})
         new_batchnorm2d = BaseNode(name=node.name + '_into_BatchNorm2d',
-                                   framework_attr={NUM_FEATURES: out_channels,
-                                                   EPSILON: EPSILON_VAL,
-                                                   MOMENTUM: MOMENTUM_VAL},
+                                   framework_attr=framework_attr,
                                    input_shape=node.output_shape,
                                    output_shape=node.output_shape,
                                    weights=bn_node_weights,

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

@@ -21,6 +21,7 @@ from typing import Dict, Tuple, List
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode, Graph
+from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.pytorch.constants import *
 from model_compression_toolkit.logger import Logger
 
@@ -38,9 +39,11 @@ class FunctionalLayerNorm(common.BaseSubstitution):
         super().__init__(matcher_instance=ln_node)
 
     @staticmethod
-    def get_attributes_from_weights(node: BaseNode, normalized_shape: [Tuple, List, int]) -> Dict:
+    def get_attributes_from_weights(node: FunctionalNode, normalized_shape: [Tuple, List, int]) -> Dict:
         """
-        Parse layer_norm(input, normalized_shape, weight=None, bias=None)
+        Convert functional layer_norm positional weights to LayerNorm weights. Extract indices of gamma
+        and beta according to tensor_input_allocs if they were input as kwargs. If they were input as args,
+        use their fixed positions.
         Args:
             node: Node that match the pattern in the substitution init.
             normalized_shape: nn.LayerNorm "normalized_shape" argument
@@ -50,28 +53,26 @@ class FunctionalLayerNorm(common.BaseSubstitution):
         """
 
         # Define default weight and bias
-        weights_dict = {GAMMA: np.ones(normalized_shape), # Default value in case weight is not given
-                        BETA: np.zeros(normalized_shape) # Default value in case bias is not given
+        weights_dict = {GAMMA: np.ones(normalized_shape),  # Default value in case weight is not given
+                        BETA: np.zeros(normalized_shape)  # Default value in case bias is not given
                         }
 
         # Check if weight and/or bias were not given.
-        has_weight = WEIGHT not in node.framework_attr
-        has_bias = BIAS not in node.framework_attr
+        if KERNEL in node.tensor_input_allocs:
+            weights_dict[GAMMA] = node.weights[node.tensor_input_allocs.index(KERNEL)]
+        elif KERNEL not in node.op_call_kwargs:
+            weights_dict[GAMMA] = node.weights[1]
 
-        if 1 in node.weights:
-            if has_weight:
-                weights_dict[GAMMA] = node.weights[1]
-            else:
-                weights_dict[BETA] = node.weights[1]
-        if 2 in node.weights:
-            assert has_bias
+        if BIAS in node.tensor_input_allocs:
+            weights_dict[BETA] = node.weights[node.tensor_input_allocs.index(BIAS)]
+        elif BIAS not in node.op_call_kwargs:
             weights_dict[BETA] = node.weights[2]
 
         return weights_dict
 
     def substitute(self,
                    graph: Graph,
-                   node: BaseNode) -> Graph:
+                   node: FunctionalNode) -> Graph:
         """
         Substitute functional.layer_norm and its inputs with LayerNorm.
         Args:
@@ -85,10 +86,11 @@ class FunctionalLayerNorm(common.BaseSubstitution):
 
         ln_node_weights = self.get_attributes_from_weights(node, normalized_shape)
 
+        framework_attr = {NORMALIZED_SHAPE: normalized_shape}
+        if EPSILON in node.op_call_kwargs:
+            framework_attr.update({EPSILON: node.op_call_kwargs[EPSILON]})
         new_layernorm = BaseNode(name=node.name + '_into_LayerNorm',
-                                 framework_attr={NORMALIZED_SHAPE: normalized_shape,
-                                                 EPSILON: node.framework_attr.get('eps'),
-                                                 },
+                                 framework_attr=framework_attr,
                                  input_shape=node.output_shape,
                                  output_shape=node.output_shape,
                                  weights=ln_node_weights,

model_compression_toolkit/core/pytorch/reader/graph_builders.py

@@ -13,11 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 import inspect
-from typing import Dict, List, Tuple, Callable
+from operator import getitem
+from typing import Dict, List, Tuple, Callable, Union, Any, Type
+
+import numpy as np
 import torch
 from torch.fx import GraphModule, Node
 
-from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.base_graph import OutTensor
 from model_compression_toolkit.core.common.graph.edge import Edge
@@ -28,29 +30,131 @@ from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlac
 from model_compression_toolkit.logger import Logger
 
 
-def extract_holder_weights(constant_name, node_target, model, weights, to_numpy):
+def _extract_parameters_and_buffers(module: Union[torch.nn.Module, GraphModule],
+                                    to_numpy: Callable) -> Dict[str, np.ndarray]:
     """
-    Extract layer weights and named buffers to a dictionary.
+    Extract parameters & buffers from input module to a dictionary.
     Args:
-        constant_name: name to write the parameters under, should be the node name.
-        node_target: relevant parameter name from Pytorch FX model.
-        model: Pytorch FX model.
-        weights: dictionary containing the weights of the node.
+        module: FX ot PyTorch module to extract parameters and buffers from.
+
+    Returns:
+        Dictionary containing module parameters and buffers by name.
+    """
+
+    named_parameters = {name: to_numpy(parameter) for name, parameter in module.named_parameters()}
+    named_buffers = {name: to_numpy(buffer) for name, buffer in module.named_buffers()}
+
+    return {**named_parameters, **named_buffers}
+
+
+def is_instance_first_arg(n: Node, expected_type: Union[Type, Tuple[Type]]) -> bool:
+    """
+    Check whether first argument of the node is the expected type
+    Args:
+        n: fx node.
+        expected_type: Expected 1st argument type.
+
+    Returns:
+        True is the first argument of node n is of the expected type, else return False.
+
+    """
+    return len(n.args) > 0 and isinstance(n.args[0], expected_type)
+
+
+def _build_input_alloc_and_call_args(n: Node, input_tensors_in_node_kwargs: Dict,
+                                     inputs_as_list: bool) -> Tuple[List, List]:
+    """
+    Build the tensor inputs list and op_call_args of the functional node.
+
+    Args:
+        n: fx node.
+        input_tensors_in_node_kwargs: A dictionary of node kwarg name and input fx node.
+        inputs_as_list: Is node's inputs are a list.
+
+    Returns:
+        A list of updated op_call args.
+        A list of tensor allocations in node's inputs.
+
+    """
+
+    tensor_input_alloc = []
+    op_call_args = list(n.args)
+    if inputs_as_list:
+        op_call_args.pop(0)
+    else:
+        for in_node in n.all_input_nodes:
+            # The extra for loop is used to tackle the case of the same input tensor for this node (e.g. torch.add(x, x)).
+            for i, arg in enumerate(n.args):
+                if arg == in_node:
+                    tensor_input_alloc.append(i)
+            for k, arg in input_tensors_in_node_kwargs.items():
+                if arg == in_node:
+                    tensor_input_alloc.append(k)
+
+    return op_call_args, tensor_input_alloc
+
+
+def _extract_torch_layer_data(node_module: torch.nn.Module,
+                              to_numpy: Callable) -> Tuple[Any, Dict[str, np.ndarray], Dict]:
+    """
+    Extract required data from a non-functional node to rebuild the PyTorch layer.
+
+    Args:
+        node_module: Torch layer, such as nn.Conv2d, nn.Linear, etc.
         to_numpy: Function to convert framework's tensor to a Numpy array.
 
     Returns:
-        Updated weights dictionary.
+        Node layer class.
+        A mapping between the layer's named parameters and buffers to their tensor values.
+        A framework_attr dictionary required to instantiate the node with the layer class.
+    """
+    node_type = type(node_module)
+    if not isinstance(node_module, torch.nn.Module):
+        Logger.error(f"Expected an instance of torch.nn.Module for node {node_module.name}, but got {node_type}")
+    # Extract the instance framework_attr (i.e. the arguments the class instance was initialized with). "fullargspec"
+    # is a list of the layer's attribute names, that will be used as keys of the framework_attr dictionary. We the
+    # values from the layer instance.
+    fullargspec = inspect.getfullargspec(node_type.__init__).args
+    framework_attr = {k: v for k, v in node_module.__dict__.items() if k in fullargspec}
+    # The "bias" argument doesn't appear in the node_module.__dict__, so we add it manually.
+    if hasattr(node_module, BIAS) and BIAS in fullargspec:
+        framework_attr[BIAS] = False if node_module.bias is None else True
+
+    # Extract layer weights and named buffers.
+    weights = {n: w for n, w in _extract_parameters_and_buffers(node_module, to_numpy).items() if len(w.shape) > 0}
+    return node_type, weights, framework_attr
+
+
+def _extract_input_and_output_shapes(_node: Node) -> Tuple[List, List]:
+    """
+    Extract input and output shapes of a node.
+    Args:
+        _node: fx node.
+
+    Returns:
+        Input and output shapes as lists.
     """
-    named_parameters_weights = {constant_name: to_numpy(parameter) for name, parameter in
-                                model.named_parameters() if node_target == name}
-    named_buffer_weights = {constant_name: to_numpy(parameter) for name, parameter in
-                            model.named_buffers() if node_target == name}
-    if len(named_parameters_weights) + len(named_buffer_weights) > 1:
-        Logger.critical("A single constant parameter must correspond to exactly one tensor. Found {len(named_parameters_weights) + len(named_buffer_weights)} parameters.")
+    input_shape = []
+    if _node.op != PLACEHOLDER:
+        for i, input_node in enumerate(_node.all_input_nodes):
+            tensor_meta = input_node.meta
+            if tensor_meta[TYPE] in [torch.Tensor, torch.nn.parameter.Parameter]:
+                input_shape += [list(tensor_meta[TENSOR_META].shape)]
+            elif tensor_meta[TYPE] == tuple:
+                input_shape += [list(n.shape) for n in tensor_meta[TENSOR_META]]
+            elif tensor_meta[TYPE] == int:
+                input_shape += [[1]]
+
+    if _node.meta[TYPE] == torch.Tensor:
+        output_shape = [list(_node.meta[TENSOR_META].shape)]
+    elif _node.meta[TYPE] in (list, tuple):
+        output_shape = [list(m.shape) for m in _node.meta[TENSOR_META]]
+    elif _node.meta[TYPE] == int:
+        output_shape = [[1]]
+    else:
+        output_shape = []
 
-    weights.update(named_parameters_weights)
-    weights.update(named_buffer_weights)
-    return weights
+    return input_shape, output_shape
 
 
 def nodes_builder(model: GraphModule,
@@ -67,135 +171,104 @@ def nodes_builder(model: GraphModule,
     Returns:
         A list of Graph nodes that were built from the fx GraphModule nodes.
     """
-    # init function variables:
-    inputs = []
-    outputs = []
-    nodes = []
-    output_nodes = []
+    # Init function variables:
+    inputs, outputs = [], []
+    nodes, output_nodes = [], []
     fx_node_2_graph_node = {}
     consts_dict = {}
     used_consts = set()
 
+    # Init parameters & buffers dictionary of the entire model. We later extract the constants values from this dictionary.
+    model_parameters_and_buffers = _extract_parameters_and_buffers(model, to_numpy)
+
     for node in model.graph.nodes:
-        # extract node type and framework attributes
-        framework_attr = dict(node.kwargs)
+
+        # ##############################################
+        #  Extract node type and framework attributes  #
+        # ##############################################
+        weights = {}
+        framework_attr = {}
         node_has_activation = True
+
         if node.target in module_dict.keys():
-            node_module = module_dict[node.target]
-            node_type = type(node_module)
-            framework_attr = node_module.__dict__
-            fullargspec = inspect.getfullargspec(node_type.__init__).args
-            framework_attr = {k: v for k, v in framework_attr.items() if k in fullargspec}
-            if hasattr(node_module, BIAS) and BIAS in fullargspec:
-                framework_attr[BIAS] = False if node_module.bias is None else True
+            # PyTorch module node, such as nn.Conv2d or nn.Linear.
+            node_type, weights, framework_attr = _extract_torch_layer_data(module_dict[node.target], to_numpy)
+
         elif node.op == CALL_FUNCTION:
+            # Node is a function that handle a parameter\buffer in the model.
             node_type = node.target
-            if node_type == getattr:
+            if node_type in [getattr, getitem]:
                 node_has_activation = False
-                Logger.warning(
-                    'Pytorch model has a parameter or constant Tensor value. This can cause unexpected behaviour when '
-                    'converting the model.')
+
         elif node.op == PLACEHOLDER:
+            # Input node to the model.
             node_type = DummyPlaceHolder
+
         elif node.op == OUTPUT:
+            # Output node of the model. Only saved in output_nodes for later handling.
             output_nodes += node.all_input_nodes
             continue
+
         elif node.op == CALL_METHOD:
+            # Node is a PyTorch function such as torch.add, torch.reshape etc.
             if hasattr(torch, node.target):
                 node_type = getattr(torch, node.target)
             elif hasattr(torch.Tensor, node.target):
                 node_type = getattr(torch.Tensor, node.target)
             else:
-                Logger.critical(f"The call method '{node.target}' is not supported.")
-        elif node.op == GET_ATTR:
-            Logger.warning(
-                'Pytorch model has a parameter or constant Tensor value. This can cause unexpected behaviour when '
-                'converting the model.')
-        else:
-            Logger.critical(f'Encountered an unsupported node type in node: {node.name}.')
+                Logger.critical(f"The call method '{node.target}' in {node} is not supported.")
 
-        # extract layer weights and named buffers
-        weights = {}
-        if node.target in module_dict.keys():
-            named_parameters_weights = {name: to_numpy(parameter) for name, parameter in
-                                        module_dict[node.target].named_parameters()}
-            named_buffer_weights = {name: to_numpy(parameter) for name, parameter in
-                                    module_dict[node.target].named_buffers() if len(parameter.shape) > 0}
-            weights.update(named_parameters_weights)
-            weights.update(named_buffer_weights)
-
-        if node.op == GET_ATTR:
-            new_const = extract_holder_weights(node, node.target, model, weights, to_numpy)
-            if list(new_const.keys())[0] in consts_dict:
+        elif node.op == GET_ATTR:
+            # Node holding a constant -> add to consts_dict so can add them later to weights of next node.
+            if node.target in consts_dict:
                 Logger.critical('A constant weight appears to have been recorded multiple times.')
-            consts_dict.update(new_const)
+            consts_dict[node] = model_parameters_and_buffers[node.target]
             continue
+        else:
+            Logger.critical(f'Encountered an unsupported node type in node: {node.name}.')
 
-        # extract input shapes and const weights
-        input_shape = []
+        # Add constants to weights dictionary.
         if node.op != PLACEHOLDER:
             for i, input_node in enumerate(node.all_input_nodes):
                 if input_node in consts_dict:
                     used_consts.add(input_node)
                     weights.update({i: consts_dict[input_node]})
 
-                tensor_meta = input_node.meta
-                if tensor_meta[TYPE] in [torch.Tensor, torch.nn.parameter.Parameter]:
-                    input_shape += [list(tensor_meta[TENSOR_META].shape)]
-                elif tensor_meta[TYPE] == tuple:
-                    input_shape += [list(n.shape) for n in tensor_meta[TENSOR_META]]
-                elif tensor_meta[TYPE] == int:
-                    input_shape += [[1]]
-
-        # extract output shapes
-        if node.meta[TYPE] == torch.Tensor:
-            output_shape = [list(node.meta[TENSOR_META].shape)]
-        elif node.meta[TYPE] in (list, tuple):
-            output_shape = [list(m.shape) for m in node.meta[TENSOR_META]]
-        elif node.meta[TYPE] == int:
-            output_shape = [[1]]
-        else:
-            output_shape = []
-
-        # filter Nodes from framework attributes, we replace these attributes with nx graph nodes
-        framework_attr_filtered = {}
-        framework_attr_nodes = {}
-        for k, v in framework_attr.items():
-            if isinstance(v, torch.fx.node.Node):
-                framework_attr_nodes[k] = v
-            else:
-                framework_attr_filtered[k] = v
-        framework_attr = framework_attr_filtered
-
-        # filter Nodes from node kwargs, we replace these attributes with nx graph nodes
-        node_kwargs = {}
-        for k, v in node.kwargs.items():
-            if not isinstance(v, torch.fx.node.Node):
-                node_kwargs[k] = v
+        # Extract input and output shapes of the node.
+        input_shape, output_shape = _extract_input_and_output_shapes(node)
 
-        # initiate graph nodes
+        # Initiate graph nodes.
         if node.op in [CALL_METHOD, CALL_FUNCTION]:
             graph_node_type = FunctionalNode
-            inputs_as_list1 = len(node.args) > 0 and isinstance(node.args[0], (list, tuple)) and all(
-                [isinstance(n, torch.fx.node.Node) for n in node.args[0]])
-            inputs_as_list = inputs_as_list1 or (len(node.args) > 0 and isinstance(node.args[0], Node) and
-                                                 node.args[0].op == PLACEHOLDER and node.args[0].meta[TYPE] in (list, tuple))
-            tensor_input_alloc = []
-            op_call_args = list(node.args)
-            if inputs_as_list:
-                op_call_args.pop(0)
-            else:
-                for in_node in node.all_input_nodes:
-                    for i, arg in enumerate(node.args):
-                        if arg == in_node:
-                            tensor_input_alloc.append(i)
-                    for k, arg in framework_attr_nodes.items():
-                        if arg == in_node:
-                            tensor_input_alloc.append(k)
-
-            # remove torch.fx.node.Node from inputs to graph_node_type
+
+            # Filter FX nodes from node_kwargs. These FX nodes are tensor inputs to the node that are part of the
+            # model's graph. We remove them because the node_kwargs should not include input tensors of the node.
+            # These input tensors will be inserted in the kwargs according to the tensor_input_alloc which is used
+            # to convert the input_tensors list in the builder to the node's args & kwargs.
+            node_kwargs, input_tensors_in_node_kwargs = {}, {}
+            for k, v in node.kwargs.items():
+                if isinstance(v, Node):
+                    input_tensors_in_node_kwargs[k] = v
+                else:
+                    node_kwargs[k] = v
+
+            # Check if node's first input argument is a list of input fx nodes, such as torch.cat:
+            is_first_input_list_of_nodes = is_instance_first_arg(node, (list, tuple)) and all(
+                [isinstance(n, Node) for n in node.args[0]])
+            is_placeholder_a_list = is_instance_first_arg(node, Node) and \
+                     node.args[0].op == PLACEHOLDER and node.args[0].meta[TYPE] in (list, tuple)
+            inputs_as_list = is_first_input_list_of_nodes or is_placeholder_a_list
+
+            # Build tensor_input_alloc required for the model builder. All input nodes are received as a list in the builder,
+            # so tensor_input_alloc is used to allocate each input tensor in the correct place in the node's args & kwargs.
+            op_call_args, tensor_input_alloc = _build_input_alloc_and_call_args(node, input_tensors_in_node_kwargs,
+                                                                                inputs_as_list)
+
+            # Remove torch.fx.node.Node from inputs to the functional node. FX nodes are input tensors in the builder,
+            # so they are remove from the op_call_args (same as op_call_kwargs) and are inserted back according to the
+            # tensor_input_alloc list.
             op_call_args = [arg for arg in op_call_args if not isinstance(arg, Node)]
-            # convert torch.fx.immutable_collections.immutable_list to tuple
+            # Convert torch.fx.immutable_collections.immutable_list to tuple.
             op_call_args = [tuple(arg) if isinstance(arg, torch.fx.immutable_collections.immutable_list) else arg
                             for arg in op_call_args]
 
@@ -205,8 +278,12 @@ def nodes_builder(model: GraphModule,
                       INPUTS_AS_LIST: inputs_as_list,
                       TENSOR_INPUT_ALLOCS: tensor_input_alloc}
         else:
+            if not all([not isinstance(v, Node) for v in framework_attr.values()]):
+                Logger.critical(f'Found FX nodes in framework attributes of {node.name}. This node type should not contain any.')  # pragma: no cover
+
             graph_node_type = BaseNode
             kwargs = {}
+
         graph_node = graph_node_type(name=node.name,
                                      framework_attr=framework_attr,
                                      input_shape=input_shape,
@@ -216,7 +293,7 @@ def nodes_builder(model: GraphModule,
                                      has_activation=node_has_activation,
                                      **kwargs)
 
-        # generate graph inputs list
+        # Generate graph inputs list.
         if node.op == PLACEHOLDER:
             for ii in range(len(output_shape)):
                 inputs.append(graph_node)
@@ -224,12 +301,12 @@ def nodes_builder(model: GraphModule,
         fx_node_2_graph_node[node] = graph_node
         nodes.append(graph_node)
 
-    # make sure all extracted constants were used in the graph
+    # Check whether all extracted constants were used in the graph.
     not_connected_consts = [c for c in consts_dict if c not in used_consts]
     if not_connected_consts:
-        Logger.critical(f'Error reading graph: These constants are not connected in the graph: {not_connected_consts}.')
+        Logger.critical(f'Error reading graph: These constants are not connected in the graph: {not_connected_consts}.')  # pragma: no cover
 
-    # generate graph outputs list
+    # Generate graph outputs list.
     for node in output_nodes:
         outputs.append(OutTensor(fx_node_2_graph_node[node], output_nodes.index(node)))
 

model_compression_toolkit/core/runner.py

@@ -216,7 +216,7 @@ def _set_final_resource_utilization(graph: Graph,
                 # No relevant nodes have been quantized with affect on the given target - since we only consider
                 # in the model's final size the quantized layers size, this means that the final size for this target
                 # is zero.
-                Logger.warning(f"No relevant quantized layers for the ru target {ru_target} were found, the recorded"
+                Logger.warning(f"No relevant quantized layers for the ru target {ru_target} were found, the recorded "
                                f"final ru for this target would be 0.")
                 final_ru_dict[ru_target] = 0
 

model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/operations_to_layers.py

@@ -148,6 +148,6 @@ class OperationsToLayers:
                 qco_by_opset_name = _current_tpc.get().tp_model.get_config_options_by_operators_set(ops2layers.name)
                 if layer in existing_layers:
                     Logger.critical(f'Found layer {layer.__name__} in more than one '
-                                 f'OperatorsSet')  # pragma: no cover
+                                    f'OperatorsSet')  # pragma: no cover
                 else:
                     existing_layers.update({layer: qco_by_opset_name})