mct-nightly 2.1.0.20240623.439__tar.gz → 2.1.0.20240624.520__tar.gz

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/PKG-INFO

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

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/mct_nightly.egg-info/PKG-INFO

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

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/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.20240623.000439"
+__version__ = "2.1.0.20240624.000520"

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/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,

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/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,

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/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:]})

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/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,

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/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,

mct-nightly-2.1.0.20240624.520/model_compression_toolkit/core/pytorch/reader/graph_builders.py

@@ -0,0 +1,353 @@
+# Copyright 2022 Sony Semiconductor 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.
+# ==============================================================================
+import inspect
+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.common import BaseNode
+from model_compression_toolkit.core.common.graph.base_graph import OutTensor
+from model_compression_toolkit.core.common.graph.edge import Edge
+from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
+from model_compression_toolkit.core.pytorch.constants import OUTPUT, PLACEHOLDER, TENSOR_META, CALL_FUNCTION, TYPE, \
+    CALL_METHOD, BIAS, FUNCTIONAL_OP, OP_CALL_KWARGS, OP_CALL_ARGS, INPUTS_AS_LIST, TENSOR_INPUT_ALLOCS, GET_ATTR
+from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlaceHolder
+from model_compression_toolkit.logger import Logger
+
+
+def _extract_parameters_and_buffers(module: Union[torch.nn.Module, GraphModule],
+                                    to_numpy: Callable) -> Dict[str, np.ndarray]:
+    """
+    Extract parameters & buffers from input module to a dictionary.
+    Args:
+        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:
+        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.
+    """
+    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 = []
+
+    return input_shape, output_shape
+
+
+def nodes_builder(model: GraphModule,
+                  module_dict: Dict,
+                  to_numpy: Callable) -> Tuple[List, List, List, Dict]:
+    """
+    Build a node from a fx node. A node contains all information to reconstruct the model module or call function
+    it's representing in the model: operation, module configuration, weights, input/output shape.
+    Args:
+        model: Pytorch FX model.
+        module_dict: A dictionary of the Pyotrch model's named modules.
+        to_numpy: A function to convert a Tensor to numpy array
+
+    Returns:
+        A list of Graph nodes that were built from the fx GraphModule 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  #
+        # ##############################################
+        weights = {}
+        framework_attr = {}
+        node_has_activation = True
+
+        if node.target in module_dict.keys():
+            # 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 in [getattr, getitem]:
+                node_has_activation = False
+
+        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}' in {node} is not supported.")
+
+        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[node] = model_parameters_and_buffers[node.target]
+            continue
+        else:
+            Logger.critical(f'Encountered an unsupported node type in node: {node.name}.')
+
+        # 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]})
+
+        # Extract input and output shapes of the node.
+        input_shape, output_shape = _extract_input_and_output_shapes(node)
+
+        # Initiate graph nodes.
+        if node.op in [CALL_METHOD, CALL_FUNCTION]:
+            graph_node_type = FunctionalNode
+
+            # 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.
+            op_call_args = [tuple(arg) if isinstance(arg, torch.fx.immutable_collections.immutable_list) else arg
+                            for arg in op_call_args]
+
+            kwargs = {FUNCTIONAL_OP: node_type,
+                      OP_CALL_ARGS: op_call_args,
+                      OP_CALL_KWARGS: node_kwargs,
+                      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,
+                                     output_shape=output_shape,
+                                     weights=weights,
+                                     layer_class=node_type,
+                                     has_activation=node_has_activation,
+                                     **kwargs)
+
+        # Generate graph inputs list.
+        if node.op == PLACEHOLDER:
+            for ii in range(len(output_shape)):
+                inputs.append(graph_node)
+
+        fx_node_2_graph_node[node] = graph_node
+        nodes.append(graph_node)
+
+    # 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}.')  # pragma: no cover
+
+    # Generate graph outputs list.
+    for node in output_nodes:
+        outputs.append(OutTensor(fx_node_2_graph_node[node], output_nodes.index(node)))
+
+    return nodes, inputs, outputs, fx_node_2_graph_node
+
+
+def edges_builder(model: GraphModule,
+                  fx_node_2_graph_node: Dict) -> List:
+    """
+
+    Args:
+        model: Pytorch FX model.
+        fx_node_2_graph_node: dictionary from fx node to graph node.
+
+    Returns:
+        List of graph edges
+    """
+    src_index = 0  # in fx src_index is always zero because fx uses the getitem operator to fetch node outputs
+    edges = []
+    connectivity_dict = {}
+    for node in model.graph.nodes:
+        if node.op != OUTPUT:
+            for input_node in node.all_input_nodes:
+                if input_node in fx_node_2_graph_node:
+                    # n_edges_for_input_node is for the case that the input node appears more than
+                    # once as the input of the node, for example add(x, x)
+                    n_edges_for_input_node = sum([1 for a in node.args if input_node == a])
+                    n_edges_for_input_node = max(n_edges_for_input_node, 1)
+
+                    dst_index = node.all_input_nodes.index(input_node)
+                    for i in range(n_edges_for_input_node):
+                        if connectivity_dict.get(input_node):
+                            connectivity_dict[input_node].append((node, dst_index))
+                        else:
+                            connectivity_dict[input_node] = [(node, dst_index)]
+                        dst_index += 1
+    for node in model.graph.nodes:
+        out_nodes = connectivity_dict.get(node)
+        if out_nodes:
+            for (out_node, dst_index) in out_nodes:
+                edges.append(
+                    Edge(fx_node_2_graph_node[node], fx_node_2_graph_node[out_node], src_index, dst_index))
+
+    return edges

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/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
 

{mct-nightly-2.1.0.20240623.439 → mct-nightly-2.1.0.20240624.520}/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})