mct-nightly 2.1.0.20240616.65727__tar.gz → 2.1.0.20240617.451__tar.gz

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240617.451}/PKG-INFO

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

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240617.451}/mct_nightly.egg-info/PKG-INFO

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

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240617.451}/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.20240616.065727"
+__version__ = "2.1.0.20240617.000451"

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240617.451}/model_compression_toolkit/core/common/graph/functional_node.py

@@ -25,7 +25,7 @@ class FunctionalNode(BaseNode):
                  functional_op: Any = None,
                  inputs_as_list: bool = False,
                  has_activation: bool = True,
-                 tensor_input_indices = None):
+                 tensor_input_allocs = None):
         """
         Init a FunctionalNode object.
 
@@ -44,7 +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_indices: A list of indices for activation tensors in the node's input tensor list
+            tensor_input_allocs: A list of indices for activation tensors in the node's input tensor list
 
         """
 
@@ -63,7 +63,7 @@ class FunctionalNode(BaseNode):
         self.op_call_args = op_call_args
         self.functional_op = functional_op
         self.inputs_as_list = inputs_as_list
-        self.tensor_input_indices = [] if tensor_input_indices is None else tensor_input_indices
+        self.tensor_input_allocs = [] if tensor_input_allocs is None else tensor_input_allocs
 
     @property
     def type(self):

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240617.451}/model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -22,6 +22,7 @@ from networkx import topological_sort
 
 from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.back2framework.base_model_builder import BaseModelBuilder
 from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX
@@ -66,8 +67,8 @@ def _build_input_tensors_list(node: BaseNode,
     return input_tensors
 
 
-def _merge_inputs(_node: BaseNode, input_tensors: List, op_call_args: List,
-                  tensor_input_indices: List = None) -> List:
+def _merge_inputs(_node: BaseNode, input_tensors: List, op_call_args: List, op_call_kwargs: Dict,
+                  tensor_input_allocs: List = None) -> Tuple[List, Dict]:
     """
     Merge input tensors list with positional weights and op_call_args, according to correct order.
 
@@ -75,22 +76,30 @@ def _merge_inputs(_node: BaseNode, input_tensors: List, op_call_args: List,
         _node: The node the inputs are for.
         input_tensors: activation input tensors to node.
         op_call_args: framework node call args.
+        op_call_kwargs: framework node call kwargs.
+        tensor_input_allocs: List of input allocations to node.
 
     Returns:
         Combined list of input_tensors and op_call_args.
     """
-    if isinstance(_node, FunctionalNode) and _node.tensor_input_indices:
+    if isinstance(_node, FunctionalNode) and _node.tensor_input_allocs:
         _input_list = op_call_args.copy()
-        if tensor_input_indices is None:
-            tensor_input_indices = _node.tensor_input_indices
-        assert len(tensor_input_indices) == len(input_tensors), \
-            f'Mismatch between input tensors ({len(tensor_input_indices)}) and indices {len(input_tensors)}'
-        for i, t in zip(tensor_input_indices, input_tensors):
-            _input_list.insert(i, t)
+        if tensor_input_allocs is None:
+            tensor_input_allocs = _node.tensor_input_allocs
+        if len(tensor_input_allocs) != len(input_tensors):
+            Logger.error(f'Mismatch between input tensors ({len(tensor_input_allocs)}) '
+                         f'and indices {len(input_tensors)} in node {_node.name}.')  # pragma: no cover
+        for i, t in zip(tensor_input_allocs, input_tensors):
+            # insert input tensors in either args or kwargs, according to tensor_input_allocs
+            if isinstance(i, str):
+                assert i not in op_call_kwargs
+                op_call_kwargs.update({i: t})
+            else:
+                _input_list.insert(i, t)
     else:
         _input_list = input_tensors + op_call_args
 
-    return _input_list
+    return _input_list, op_call_kwargs
 
 
 def _run_operation(n: BaseNode,
@@ -125,14 +134,15 @@ def _run_operation(n: BaseNode,
         # 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_indices = None
+        _tensor_input_allocs = None
     else:
-        _tensor_input_indices = [i for i in n.tensor_input_indices if i not in n.weights]
+        _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)
     else:
-        merged_inputs = _merge_inputs(n, input_tensors, op_call_args, tensor_input_indices=_tensor_input_indices)
+        merged_inputs, functional_kwargs = _merge_inputs(n, input_tensors, op_call_args, functional_kwargs.copy(),
+                                                         tensor_input_allocs=_tensor_input_allocs)
         out_tensors_of_n_float = op_func(*merged_inputs, **functional_kwargs)
 
     # Add a fake quant node if the node has an activation threshold.

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240617.451}/model_compression_toolkit/core/pytorch/constants.py

@@ -40,7 +40,7 @@ FUNCTIONAL_OP = 'functional_op'
 OP_CALL_ARGS = 'op_call_args'
 OP_CALL_KWARGS = 'op_call_kwargs'
 INPUTS_AS_LIST = 'inputs_as_list'
-TENSOR_INPUT_INDICES = 'tensor_input_indices'
+TENSOR_INPUT_ALLOCS = 'tensor_input_allocs'
 INPLACE = 'inplace'
 HARDTANH_MIN_VAL = 'min_val'
 HARDTANH_MAX_VAL = 'max_val'

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240617.451}/model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/reshape_with_static_shapes.py

@@ -65,10 +65,10 @@ class ReshapeWithStaticShapes(common.BaseSubstitution):
 
         # When a "reshape" is called with multiple arguments (e.g. x.reshape(-1, channels, height, width)
         # this substitution converts it x.reshape((-1, channels, height, width)), so need to update the
-        # tensor_input_indices attribute.
-        # scalar argument's shape is [1] so remove those indices from tensor_input_indices
+        # tensor_input_allocs attribute.
+        # scalar argument's shape is [1] so remove those indices from tensor_input_allocs
         # node.input_shape example: [[1, 32, 4, 32], [1], [1], [1]]
-        node.tensor_input_indices = node.tensor_input_indices[:sum([i != [1] for i in node.input_shape])]
+        node.tensor_input_allocs = node.tensor_input_allocs[:sum([i != [1] for i in node.input_shape])]
 
         # modify the node input info
         node.input_shape = [node.input_shape[0]]

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240617.451}/model_compression_toolkit/core/pytorch/reader/graph_builders.py

@@ -23,7 +23,7 @@ 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_INDICES, GET_ATTR
+    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
 
@@ -140,7 +140,7 @@ def nodes_builder(model: GraphModule,
                     weights.update({i: consts_dict[input_node]})
 
                 tensor_meta = input_node.meta
-                if tensor_meta[TYPE] == torch.Tensor:
+                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]]
@@ -159,8 +159,11 @@ def nodes_builder(model: GraphModule,
 
         # 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 not isinstance(v, torch.fx.node.Node):
+            if isinstance(v, torch.fx.node.Node):
+                framework_attr_nodes[k] = v
+            else:
                 framework_attr_filtered[k] = v
         framework_attr = framework_attr_filtered
 
@@ -177,7 +180,7 @@ def nodes_builder(model: GraphModule,
                 [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_index = []
+            tensor_input_alloc = []
             op_call_args = list(node.args)
             if inputs_as_list:
                 op_call_args.pop(0)
@@ -185,7 +188,10 @@ def nodes_builder(model: GraphModule,
                 for in_node in node.all_input_nodes:
                     for i, arg in enumerate(node.args):
                         if arg == in_node:
-                            tensor_input_index.append(i)
+                            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
             op_call_args = [arg for arg in op_call_args if not isinstance(arg, Node)]
@@ -197,7 +203,7 @@ def nodes_builder(model: GraphModule,
                       OP_CALL_ARGS: op_call_args,
                       OP_CALL_KWARGS: node_kwargs,
                       INPUTS_AS_LIST: inputs_as_list,
-                      TENSOR_INPUT_INDICES: tensor_input_index}
+                      TENSOR_INPUT_ALLOCS: tensor_input_alloc}
         else:
             graph_node_type = BaseNode
             kwargs = {}