mct-nightly 2.2.0.20240924.500__tar.gz → 2.2.0.20240926.452__tar.gz

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/PKG-INFO

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

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/mct_nightly.egg-info/PKG-INFO

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

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/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.20240924.000500"
+__version__ = "2.2.0.20240926.000452"

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/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)

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/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)

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/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)

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/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}"

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/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,

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/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.

{mct-nightly-2.2.0.20240924.500 → mct-nightly-2.2.0.20240926.452}/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])