mct-nightly 2.2.0.20240925.453__py3-none-any.whl → 2.2.0.20240926.452__py3-none-any.whl

{mct_nightly-2.2.0.20240925.453.dist-info → mct_nightly-2.2.0.20240926.452.dist-info}/METADATA

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

{mct_nightly-2.2.0.20240925.453.dist-info → mct_nightly-2.2.0.20240926.452.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=eI5QONSxEPDPEpc5TqZF43AfSWl8Om61_pe6SgWyCTk,1573
+model_compression_toolkit/__init__.py,sha256=vX11_K5A8c4_uT3X2dHKRg0nxBh-qKBSqljT0u_1B64,1573
 model_compression_toolkit/constants.py,sha256=i4wYheBkIdQmsQA-axIpcT3YiSO1USNc-jaNiNE8w6E,3920
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -47,7 +47,7 @@ model_compression_toolkit/core/common/graph/memory_graph/max_cut_astar.py,sha256
 model_compression_toolkit/core/common/graph/memory_graph/memory_element.py,sha256=gRmBEFRmyJsNKezQfiwDwQu1cmbGd2wgKCRTH6iw8mw,3961
 model_compression_toolkit/core/common/graph/memory_graph/memory_graph.py,sha256=gw4av_rzn_3oEAPpD3B7PHZDqnxHMjIESevl6ppPnkk,7175
 model_compression_toolkit/core/common/hessian/__init__.py,sha256=6216QgHl7h4DXGn5ForP9Tija-wrBSONNtQ769ikP2s,1025
-model_compression_toolkit/core/common/hessian/hessian_info_service.py,sha256=DHbZqFDuDir1QWN-YkYBzaoGDujgYam1hT2ea6uL3yM,21009
+model_compression_toolkit/core/common/hessian/hessian_info_service.py,sha256=TfgSIh5pmZcJM9335aAxZriCzMljnk3mYhmKBsK2x5Y,20848
 model_compression_toolkit/core/common/hessian/hessian_info_utils.py,sha256=1axmN0tjJSo_7hUr2d2KMv4y1pBi19cqWSQpi4BbdsA,1458
 model_compression_toolkit/core/common/hessian/hessian_scores_calculator.py,sha256=Pe4uKerx-MeDQPJ7Slr8fvFUHfv02q33w3gbQK5kBKs,4186
 model_compression_toolkit/core/common/hessian/hessian_scores_request.py,sha256=atGJgJBL9uwYRC3t9NnzGgHYxV4XJj4Ai_xPpQH0rhY,3229
@@ -228,7 +228,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=QfMulTLC6X_0Iwvk_VChFIoSdeoiEeJ_rf2IQi5TjBk,19353
+model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py,sha256=BJeKGMv5VU4Z3jLOIQ-Ifs_2vGELQSmEQmje3ZmaUl4,19948
 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
@@ -239,7 +239,7 @@ 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=iTuP1hjuTZTGcE7izfs_UOWBGeEBFRvRIU4QCh-b21M,4627
+model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/const_holder_conv.py,sha256=sw3jIOUSvfWUeD8l3rGcUOtC6QuzpMIQm8V3RQAM53Q,4741
 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
@@ -267,7 +267,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=BvBj9uokKTvX-6d39yA4SKwRQAN8_X4T8l-rPibChJQ,16754
+model_compression_toolkit/core/pytorch/reader/graph_builders.py,sha256=mo1NIYXxiAigbTZvNgQeLi6vzLn0RqU0RxcxZKE27cE,19335
 model_compression_toolkit/core/pytorch/reader/node_holders.py,sha256=7XNc7-l1MZPJGcOESvtAwfIMxrU6kvt3YjF5B7qOqK4,1048
 model_compression_toolkit/core/pytorch/reader/reader.py,sha256=GEJE0QX8XJFWbYCkbRBtzttZtmmuoACLx8gw9KyAQCE,6015
 model_compression_toolkit/core/pytorch/statistics_correction/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
@@ -301,7 +301,7 @@ model_compression_toolkit/data_generation/pytorch/image_operations.py,sha256=KUQ
 model_compression_toolkit/data_generation/pytorch/image_pipeline.py,sha256=dcQr-67u9-ggGuS39YAvR7z-Y0NOdJintcVQ5vy1bM8,7478
 model_compression_toolkit/data_generation/pytorch/model_info_exctractors.py,sha256=y6vMed6lQQj67-BXZKrAcWUNTkH8YjiUhknOV4wSpRA,9399
 model_compression_toolkit/data_generation/pytorch/optimization_utils.py,sha256=vRMeUEdInPuJisiO-SKo_9miWZV90sz8GCg5MY0AqiU,18098
-model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py,sha256=cUkFg-9LWwRKy11tlASJwp1FbDx6a7sZWpJNMz01hWA,21626
+model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py,sha256=_BFy4RYcLoxpt5KecM5VbPRRNM4QHdFr9WmtL4FODUE,21796
 model_compression_toolkit/data_generation/pytorch/optimization_functions/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/data_generation/pytorch/optimization_functions/batchnorm_alignment_functions.py,sha256=dMc4zz9XfYfAT4Cxns57VgvGZWPAMfaGlWLFyCyl8TA,1968
 model_compression_toolkit/data_generation/pytorch/optimization_functions/bn_layer_weighting_functions.py,sha256=We0fVMQ4oU7Y0IWQ8fKy8KpqkIiLyKoQeF9XKAQ6TH0,3317
@@ -472,9 +472,9 @@ model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tpc_keras.py,sha256=XM6qBLIvzsmdFf-AZq5WOlORK2GXC_X-gulReNxHb9E,6601
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tpc_pytorch.py,sha256=nP05jqvh6uaj30a3W7zEkJfKtqfP0Nz5bobwRqbYrdM,5807
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/__init__.py,sha256=tHTUvsaerSfbe22pU0kIDauPpFD7Pq5EmZytVIDkHz4,717
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tp_model.py,sha256=_LHQkGB0x12FQBDIkEA-Br8HSUL5ZmMXxI7lDpVWcQU,15422
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tp_model.py,sha256=Ee7M3YVymdv6HYsm7coB8N0dyTOhlAhLdxfSLJXCuoU,15665
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_keras.py,sha256=u8qD1XkHwU4LIoNbmC5mtZd8lZ8gZ4XFihZmoYwAulc,7641
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_pytorch.py,sha256=EaQERA7XnZcF7pO4xzBk0li96JnACRE7ppgK535EMXM,6698
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_pytorch.py,sha256=GCghKkkZOKNTAzwyoZZPid9alGiufNUBzDj2yE7YUSU,6709
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/target_platform_capabilities.py,sha256=is00rNrDmmirYsyMtMkWz0DwOA92-x7hAJwpd6z1n2E,2806
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/latest/__init__.py,sha256=CXC-HQolSDu7j8V-Xm-SWGCd74gXB3XnAkEhI_TVbIQ,1516
@@ -550,8 +550,8 @@ tests_pytest/pytorch/gptq/test_annealing_cfg.py,sha256=hGC7L6mp3N1ygcJ3OctgS_Fz2
 tests_pytest/pytorch/gptq/test_gradual_act_quantization.py,sha256=tI01aFIUaiCILL5Qn--p1E_rLBUelxLdSY3k52lwcx0,4594
 tests_pytest/pytorch/trainable_infrastructure/__init__.py,sha256=RAe8mgIr1V8dRIQtLf_dSG5zTUCKuQzxyybYx1dzEAs,697
 tests_pytest/pytorch/trainable_infrastructure/test_linear_annealing.py,sha256=eNOpSp0GoLxtEdiRypBp8jaujXfdNxBwKh5Rd-P7WLs,1786
-mct_nightly-2.2.0.20240925.453.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20240925.453.dist-info/METADATA,sha256=-FYOCuanQ2MY1g5nj-LWf1WuTJEAqIpR5ymgLPAGe2I,20813
-mct_nightly-2.2.0.20240925.453.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-mct_nightly-2.2.0.20240925.453.dist-info/top_level.txt,sha256=csdfSXhtRnpWYRzjZ-dRLIhOmM2TEdVXUxG05A5fgb8,39
-mct_nightly-2.2.0.20240925.453.dist-info/RECORD,,
+mct_nightly-2.2.0.20240926.452.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20240926.452.dist-info/METADATA,sha256=AWRxoKCjgmTLCf726oR4aaHZBpcxel0TrFkAGP-5guM,20813
+mct_nightly-2.2.0.20240926.452.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+mct_nightly-2.2.0.20240926.452.dist-info/top_level.txt,sha256=csdfSXhtRnpWYRzjZ-dRLIhOmM2TEdVXUxG05A5fgb8,39
+mct_nightly-2.2.0.20240926.452.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.2.0.20240925.000453"
+__version__ = "2.2.0.20240926.000452"

model_compression_toolkit/core/common/hessian/hessian_info_service.py

@@ -258,12 +258,10 @@ class HessianInfoService:
                     f"{hessian_scores_request.target_nodes}.")
 
         # Replace node in reused target nodes with a representing node from the 'reuse group'.
-        for n in hessian_scores_request.target_nodes:
-            if n.reuse_group:
-                rep_node = self._get_representing_of_reuse_group(n)
-                hessian_scores_request.target_nodes.remove(n)
-                if rep_node not in hessian_scores_request.target_nodes:
-                    hessian_scores_request.target_nodes.append(rep_node)
+        hessian_scores_request.target_nodes = [
+            self._get_representing_of_reuse_group(node) if node.reuse else node
+            for node in hessian_scores_request.target_nodes
+        ]
 
         # Ensure the saved info has the required number of approximations
         self._populate_saved_info_to_size(hessian_scores_request, required_size, batch_size)

model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -231,6 +231,7 @@ class PytorchModel(torch.nn.Module):
         self.return_float_outputs = return_float_outputs
         self.wrapper = wrapper
         self.get_activation_quantizer_holder = get_activation_quantizer_holder_fn
+        self.reuse_groups = {}
         self._add_modules()
 
     # todo: Move to parent class BaseModelBuilder
@@ -288,7 +289,19 @@ class PytorchModel(torch.nn.Module):
         Build and add the modules and functional nodes from node_sort list as attributes to PytorchModel
         """
         for node in self.node_sort:
-            node_op = self.wrap(node)
+            if node.reuse:
+                # If the node is reused, retrieve the original module
+                if node.reuse_group not in self.reuse_groups:
+                    Logger.critical(f"Reuse group {node.reuse_group} not found for node {node.name}")
+
+                node_op = self.reuse_groups[node.reuse_group]
+            else:
+                # If it's not reused, create a new module
+                node_op = self.wrap(node)
+                if node.reuse_group:
+                    # Store the module for future reuse
+                    self.reuse_groups[node.reuse_group] = node_op
+
             if isinstance(node, FunctionalNode):
                 # for functional layers
                 setattr(self, node.name, node_op)

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

@@ -80,7 +80,9 @@ class FunctionalConvSubstitution(common.BaseSubstitution):
                             output_shape=func_node.output_shape,
                             weights={KERNEL: weight} if bias is None else {KERNEL: weight, BIAS: bias},
                             layer_class=new_layer,
-                            has_activation=func_node.has_activation)
+                            has_activation=func_node.has_activation,
+                            reuse=func_node.reuse,
+                            reuse_group=func_node.reuse_group)
         graph.add_node(new_node)
         graph.reconnect_out_edges(current_node=func_node, new_node=new_node)
         graph.reconnect_in_edges(current_node=func_node, new_node=new_node)

model_compression_toolkit/core/pytorch/reader/graph_builders.py

@@ -30,8 +30,7 @@ from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlac
 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]:
+def _extract_parameters_and_buffers(module: Union[torch.nn.Module, GraphModule]) -> Dict[str, np.ndarray]:
     """
     Extract parameters & buffers from input module to a dictionary.
     Args:
@@ -41,8 +40,8 @@ def _extract_parameters_and_buffers(module: Union[torch.nn.Module, GraphModule],
         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()}
+    named_parameters = {name: parameter for name, parameter in module.named_parameters()}
+    named_buffers = {name: buffer for name, buffer in module.named_buffers()}
 
     return {**named_parameters, **named_buffers}
 
@@ -97,14 +96,12 @@ def _build_input_alloc_and_call_args(n: Node, input_tensors_in_node_kwargs: Dict
     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]:
+def _extract_torch_layer_data(node_module: torch.nn.Module) -> 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.
@@ -124,7 +121,7 @@ def _extract_torch_layer_data(node_module: torch.nn.Module,
         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}
+    weights = {n: w for n, w in _extract_parameters_and_buffers(node_module).items() if len(w.shape) > 0}
     return node_type, weights, framework_attr
 
 
@@ -181,8 +178,11 @@ def nodes_builder(model: GraphModule,
     consts_dict = {}
     used_consts = set()
 
+    # Dictionary to track seen targets and their corresponding nodes to mark reused nodes
+    seen_targets = {}
+
     # 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)
+    model_parameters_and_buffers = _extract_parameters_and_buffers(model)
 
     for node in model.graph.nodes:
 
@@ -195,7 +195,7 @@ def nodes_builder(model: GraphModule,
 
         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)
+            node_type, weights, framework_attr = _extract_torch_layer_data(module_dict[node.target])
 
         elif node.op == CALL_FUNCTION:
             # Node is a function that handle a parameter\buffer in the model.
@@ -249,6 +249,31 @@ def nodes_builder(model: GraphModule,
         # Extract input and output shapes of the node.
         input_shape, output_shape = _extract_input_and_output_shapes(node)
 
+        # Check if this node's target has been seen before
+        reuse = False
+        reuse_group = None
+        node_group_key = create_reuse_group(node.target, weights)
+        # We mark nodes as reused only if there are multiple nodes in the graph with same
+        # 'target' and it has some weights.
+        if node_group_key in seen_targets and len(weights) > 0:
+            reuse = True
+            reuse_group = node_group_key
+            # Update the 'base/main' node with the reuse group as all other nodes in its group.
+            fx_node_2_graph_node[seen_targets[node_group_key]].reuse_group = reuse_group
+        else:
+            seen_targets[node_group_key] = node
+
+        # Convert weights to numpy arrays after reuse marking
+        # We delay this conversion to preserve the original tensor instances during the reuse identification process.
+        # This is crucial for correctly identifying identical weight instances in reused functional layers.
+        # By keeping the original PyTorch tensors until this point, we ensure that:
+        # 1. Reused layers with the same weight instances are correctly marked as reused.
+        # 2. The instance-based weight signature generation works as intended, using the memory
+        # addresses of the original tensors.
+        # Only after all reuse marking is complete do we convert to numpy arrays.
+        for weight_name, weight_value in weights.items():
+            weights[weight_name] = to_numpy(weight_value)
+
         # Initiate graph nodes.
         if node.op in [CALL_METHOD, CALL_FUNCTION]:
             graph_node_type = FunctionalNode
@@ -300,6 +325,8 @@ def nodes_builder(model: GraphModule,
                                      weights=weights,
                                      layer_class=node_type,
                                      has_activation=node_has_activation,
+                                     reuse=reuse,
+                                     reuse_group=reuse_group,
                                      **kwargs)
 
         # Generate graph inputs list.
@@ -365,3 +392,28 @@ def edges_builder(model: GraphModule,
                     Edge(fx_node_2_graph_node[node], fx_node_2_graph_node[out_node], src_index, dst_index))
 
     return edges
+
+
+def create_reuse_group(target: Any, weights: Dict[str, Any]) -> str:
+    """
+    Combine target and weights to create a unique reuse group identifier.
+    We consider the weights as part of the group identifier because they are not part of
+    the module in functional layers, but if a functional layer is using the same weights multiple
+    times it is considered to be reused.
+
+    This function creates a unique string identifier for a reuse group by combining
+    the target (typically a layer or operation name) with the weights IDs.
+
+    Args:
+        target (Any): The target of the node, typically a string or callable representing
+                      a layer or operation.
+        weights (Dict[str, Any]): A dictionary of weight names to weight values.
+                                  The values can be any type (typically tensors or arrays).
+
+    Returns:
+        str: A unique string identifier for the reuse group.
+    """
+    if not weights:
+        return str(target)
+    weight_ids = tuple(sorted(id(weight) for weight in weights.values()))
+    return f"{target}_{weight_ids}"

model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+import copy
+
 import time
 from typing import Callable, Any, Tuple, List, Union
 
@@ -179,8 +181,11 @@ if FOUND_TORCH and FOUND_TORCHVISION:
         # get the model device
         device = get_working_device()
 
+        # copy model for data generation
+        model_for_data_gen = copy.deepcopy(model)
+
         # get a static graph representation of the model using torch.fx
-        fx_model = symbolic_trace(model)
+        fx_model = symbolic_trace(model_for_data_gen)
 		
 		# Get Data Generation functions and classes
         image_pipeline, normalization, bn_layer_weighting_fn, bn_alignment_loss_fn, output_loss_fn, \
@@ -208,23 +213,23 @@ if FOUND_TORCH and FOUND_TORCHVISION:
         scheduler = scheduler_get_fn(data_generation_config.n_iter)
 
         # Set the current model
-        set_model(model)
+        set_model(model_for_data_gen)
 
         # Create an activation extractor object to extract activations from the model
         activation_extractor = PytorchActivationExtractor(
-            model,
+            model_for_data_gen,
             fx_model,
             data_generation_config.bn_layer_types,
             data_generation_config.last_layer_types)
 
         # Create an orig_bn_stats_holder object to hold original BatchNorm statistics
-        orig_bn_stats_holder = PytorchOriginalBNStatsHolder(model, data_generation_config.bn_layer_types)
+        orig_bn_stats_holder = PytorchOriginalBNStatsHolder(model_for_data_gen, data_generation_config.bn_layer_types)
         if orig_bn_stats_holder.get_num_bn_layers() == 0:
             Logger.critical(
                 f'Data generation requires a model with at least one BatchNorm layer.') # pragma: no cover
 
         # Create an ImagesOptimizationHandler object for handling optimization
-        all_imgs_opt_handler = PytorchImagesOptimizationHandler(model=model,
+        all_imgs_opt_handler = PytorchImagesOptimizationHandler(model=model_for_data_gen,
                                                                 data_gen_batch_size=data_generation_config.data_gen_batch_size,
                                                                 init_dataset=init_dataset,
                                                                 optimizer=data_generation_config.optimizer,

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tp_model.py

@@ -207,7 +207,9 @@ def generate_tp_model(default_config: OpQuantizationConfig,
                                                                                   base_config=const_config_input16_per_tensor)
 
     qpreserving_const_config = const_config.clone_and_edit(enable_activation_quantization=False,
-                                                           quantization_preserving=True)
+                                                           quantization_preserving=True,
+                                                           default_weight_attr_config=const_config.default_weight_attr_config.clone_and_edit(
+                                                               weights_per_channel_threshold=False))
     qpreserving_const_config_options = tp.QuantizationConfigOptions([qpreserving_const_config])
 
     # Create a TargetPlatformModel and set its default quantization config.

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_pytorch.py

@@ -19,7 +19,7 @@ import torch
 from torch import add, sub, mul, div, divide, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, \
     chunk, unbind, topk, gather, equal, transpose, permute, argmax, squeeze, multiply, subtract
 from torch.nn import Conv2d, Linear, ConvTranspose2d, MaxPool2d
-from torch.nn import Dropout, Flatten, Hardtanh, Identity
+from torch.nn import Dropout, Flatten, Hardtanh
 from torch.nn import ReLU, ReLU6, PReLU, SiLU, Sigmoid, Tanh, Hardswish, LeakyReLU
 from torch.nn.functional import relu, relu6, prelu, silu, hardtanh, hardswish, leaky_relu
 
@@ -87,7 +87,7 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                                     squeeze,
                                                                     permute,
                                                                     transpose])
-        tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, [gather])
+        tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, [gather, torch.Tensor.expand])
         tp.OperationsSetToLayers(OPSET_MERGE_OPS,
                                  [torch.stack, torch.cat, torch.concat, torch.concatenate])