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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.1.0.20240616.65727
+Version: 2.1.0.20240618.432
 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.20240618.432}/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.20240618.432
 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.20240618.432}/mct_nightly.egg-info/SOURCES.txt

@@ -499,4 +499,30 @@ model_compression_toolkit/trainable_infrastructure/keras/load_model.py
 model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py
-model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py
+model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py
+model_compression_toolkit/xquant/__init__.py
+model_compression_toolkit/xquant/common/__init__.py
+model_compression_toolkit/xquant/common/constants.py
+model_compression_toolkit/xquant/common/core_report_generator.py
+model_compression_toolkit/xquant/common/dataset_utils.py
+model_compression_toolkit/xquant/common/framework_report_utils.py
+model_compression_toolkit/xquant/common/model_analyzer.py
+model_compression_toolkit/xquant/common/model_folding_utils.py
+model_compression_toolkit/xquant/common/similarity_calculator.py
+model_compression_toolkit/xquant/common/similarity_functions.py
+model_compression_toolkit/xquant/common/tensorboard_utils.py
+model_compression_toolkit/xquant/common/xquant_config.py
+model_compression_toolkit/xquant/keras/__init__.py
+model_compression_toolkit/xquant/keras/dataset_utils.py
+model_compression_toolkit/xquant/keras/facade_xquant_report.py
+model_compression_toolkit/xquant/keras/keras_report_utils.py
+model_compression_toolkit/xquant/keras/model_analyzer.py
+model_compression_toolkit/xquant/keras/similarity_functions.py
+model_compression_toolkit/xquant/keras/tensorboard_utils.py
+model_compression_toolkit/xquant/pytorch/__init__.py
+model_compression_toolkit/xquant/pytorch/dataset_utils.py
+model_compression_toolkit/xquant/pytorch/facade_xquant_report.py
+model_compression_toolkit/xquant/pytorch/model_analyzer.py
+model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py
+model_compression_toolkit/xquant/pytorch/similarity_functions.py
+model_compression_toolkit/xquant/pytorch/tensorboard_utils.py

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240618.432}/mct_nightly.egg-info/requires.txt

@@ -1,7 +1,7 @@
 networkx!=2.8.1
 tqdm
 Pillow
-numpy
+numpy<2.0
 opencv-python
 scikit-image
 scikit-learn

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240618.432}/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.20240618.000432"

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240618.432}/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.20240618.432}/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.20240618.432}/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.20240618.432}/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.20240618.432}/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 = {}

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/gptq/keras/gptq_training.py

@@ -353,7 +353,7 @@ class KerasGPTQTrainer(GPTQTrainer):
                     node.final_activation_quantization_cfg.set_quant_config_attr(config_attr, config_value)
                 if self.gptq_config.train_bias:
                     use_bias = layer.layer.get_config().get(USE_BIAS)
-                    if use_bias is not None and use_bias:
+                    if use_bias is not None and use_bias and layer.layer.bias is not None:
                         new_bias = layer.layer.bias.numpy()
                         node.set_weights_by_keys(BIAS, new_bias)
 

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/gptq/keras/graph_info.py

@@ -63,7 +63,7 @@ def get_gptq_trainable_parameters(fxp_model: Model,
                 kernel_ops_attrs = fw_info.kernel_ops_attributes_mapping.get(type(layer.layer))
                 use_bias = kernel_ops_attrs is not None and kernel_ops_attrs[0] is not None \
                            and layer.layer.get_config().get(USE_BIAS)
-                if use_bias is not None and use_bias:
+                if use_bias is not None and use_bias and layer.layer.bias is not None:
                     bias_weights.append([layer.layer.bias])
 
     return trainable_weights, bias_weights, trainable_threshold

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/gptq/pytorch/gptq_training.py

@@ -299,7 +299,9 @@ class PytorchGPTQTrainer(GPTQTrainer):
                 for config_attr, config_value in activation_quant_config.items():
                     node.final_activation_quantization_cfg.set_quant_config_attr(config_attr, config_value)
                 if self.gptq_config.train_bias and hasattr(layer.layer, BIAS):
-                    node.set_weights_by_keys(BIAS, self.fw_impl.to_numpy(getattr(layer.layer, BIAS)))
+                    bias = getattr(layer.layer, BIAS)
+                    if bias is not None:
+                        node.set_weights_by_keys(BIAS, self.fw_impl.to_numpy(bias))
 
         return graph_quant
 
@@ -316,4 +318,5 @@ class PytorchGPTQTrainer(GPTQTrainer):
             if isinstance(layer, PytorchQuantizationWrapper):
                 if hasattr(layer.layer, BIAS):
                     bias = getattr(layer.layer, BIAS)
-                    bias.requires_grad = self.gptq_config.train_bias
+                    if bias is not None:
+                        bias.requires_grad = self.gptq_config.train_bias

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/gptq/pytorch/graph_info.py

@@ -56,7 +56,8 @@ def get_gptq_trainable_parameters(fxp_model: nn.Module,
 
             if add_bias and hasattr(layer.layer, BIAS):
                 bias = getattr(layer.layer, BIAS)
-                trainable_bias.append(bias)
+                if bias is not None:
+                    trainable_bias.append(bias)
 
     return trainable_aux_weights, trainable_bias, trainable_threshold
 

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3/tpc_pytorch.py

@@ -18,7 +18,7 @@ import operator
 import torch
 from torch import add, sub, mul, div, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, chunk, unbind, topk, \
     gather, equal, transpose, permute, argmax, squeeze
-from torch.nn import Conv2d, Linear, ConvTranspose2d
+from torch.nn import Conv2d, Linear, ConvTranspose2d, MaxPool2d
 from torch.nn import Dropout, Flatten, Hardtanh, Identity
 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
@@ -83,7 +83,8 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                     argmax,
                                                     gather,
                                                     topk,
-                                                    squeeze])
+                                                    squeeze,
+                                                    MaxPool2d])
 
         tp.OperationsSetToLayers("Conv", [Conv2d, ConvTranspose2d],
                                  attr_mapping=pytorch_linear_attr_mapping)

{mct-nightly-2.1.0.20240616.65727 → mct-nightly-2.1.0.20240618.432}/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tpc_pytorch.py

@@ -18,7 +18,7 @@ import operator
 import torch
 from torch import add, sub, mul, div, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, chunk, unbind, topk, \
     gather, equal, transpose, permute, argmax, squeeze
-from torch.nn import Conv2d, Linear, ConvTranspose2d
+from torch.nn import Conv2d, Linear, ConvTranspose2d, MaxPool2d
 from torch.nn import Dropout, Flatten, Hardtanh, Identity
 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
@@ -82,7 +82,8 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                     argmax,
                                                     gather,
                                                     topk,
-                                                    squeeze])
+                                                    squeeze,
+                                                    MaxPool2d])
 
         tp.OperationsSetToLayers("Conv", [Conv2d, ConvTranspose2d],
                                  attr_mapping=pytorch_linear_attr_mapping)

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/__init__.py

@@ -0,0 +1,19 @@
+#  Copyright 2024 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.
+#  ==============================================================================
+
+from model_compression_toolkit.xquant.common.xquant_config import XQuantConfig
+from model_compression_toolkit.xquant.keras.facade_xquant_report import xquant_report_keras_experimental
+from model_compression_toolkit.xquant.pytorch.facade_xquant_report import xquant_report_pytorch_experimental
+

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/__init__.py

@@ -0,0 +1,15 @@
+#  Copyright 2024 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.
+#  ==============================================================================
+#

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/constants.py

@@ -0,0 +1,38 @@
+#  Copyright 2024 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.
+#  ==============================================================================
+
+# Default similarity metric names:
+CS_SIMILARITY_METRIC_NAME = 'cs'
+SQNR_SIMILARITY_METRIC_NAME = 'sqnr'
+MSE_SIMILARITY_METRIC_NAME = 'mse'
+
+# Report components names:
+OUTPUT_SIMILARITY_METRICS_REPR = 'output_similarity_metrics_repr'
+OUTPUT_SIMILARITY_METRICS_VAL = 'output_similarity_metrics_val'
+INTERMEDIATE_SIMILARITY_METRICS_REPR = 'intermediate_similarity_metrics_repr'
+INTERMEDIATE_SIMILARITY_METRICS_VAL = 'intermediate_similarity_metrics_val'
+
+# Graph attribute names:
+XQUANT_REPR = 'xquant_repr'
+XQUANT_VAL = 'xquant_val'
+
+# Report file name:
+REPORT_FILENAME = 'quant_report.json'
+
+# Tag to use in tensorboard for the graph we plot:
+TENSORBOARD_DEFAULT_TAG = 'xquant'
+
+# When extracting the activations of a model we hold the output using a dedicated key:
+MODEL_OUTPUT_KEY = 'model_output_key'

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/core_report_generator.py

@@ -0,0 +1,83 @@
+#  Copyright 2024 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.
+#  ==============================================================================
+from tqdm import tqdm
+from typing import Callable, Any, Dict
+
+from model_compression_toolkit.core.common.model_collector import ModelCollector
+from model_compression_toolkit.xquant import XQuantConfig
+from model_compression_toolkit.xquant.common.constants import OUTPUT_SIMILARITY_METRICS_REPR, OUTPUT_SIMILARITY_METRICS_VAL, INTERMEDIATE_SIMILARITY_METRICS_REPR, \
+    INTERMEDIATE_SIMILARITY_METRICS_VAL
+from model_compression_toolkit.xquant.common.framework_report_utils import FrameworkReportUtils
+
+
+def core_report_generator(float_model: Any,
+                          quantized_model: Any,
+                          repr_dataset: Callable,
+                          validation_dataset: Callable,
+                          fw_report_utils: FrameworkReportUtils,
+                          xquant_config: XQuantConfig) -> Dict[str, Any]:
+    """
+    Generate report in tensorboard with a graph of the quantized model and similarity metrics that
+    have been measured when comparing to the float model (or any other two models).
+    The report also contains histograms that are collected on the baseline model (usually, the float
+    model).
+
+    Args:
+        float_model (Any): The original floating-point model.
+        quantized_model (Any): The model after quantization.
+        repr_dataset (Callable): Representative dataset used for similarity metrics computation.
+        validation_dataset (Callable): Validation dataset used for similarity metrics computation.
+        fw_report_utils (FrameworkReportUtils): Utilities for generating framework-specific reports.
+        xquant_config (XQuantConfig): Configuration settings for explainable quantization.
+
+    Returns:
+        Dict[str, Any]: A dictionary containing the collected similarity metrics and report data.
+    """
+    # Collect histograms on the float model.
+    float_graph = fw_report_utils.model_folding_utils.create_float_folded_graph(float_model, repr_dataset)
+    mi = ModelCollector(float_graph, fw_report_utils.fw_impl, fw_report_utils.fw_info)
+    for _data in tqdm(repr_dataset(), desc="Collecting Histograms"):
+        mi.infer(_data)
+
+    # Collect histograms and add them to Tensorboard.
+    fw_report_utils.tb_utils.add_histograms_to_tensorboard(graph=float_graph)
+
+    # Compute similarity metrics on representative dataset and validation set.
+    repr_similarity = fw_report_utils.similarity_calculator.compute_similarity_metrics(float_model=float_model,
+                                                                                       quantized_model=quantized_model,
+                                                                                       dataset=repr_dataset,
+                                                                                       custom_similarity_metrics=xquant_config.custom_similarity_metrics)
+    val_similarity = fw_report_utils.similarity_calculator.compute_similarity_metrics(float_model=float_model,
+                                                                                      quantized_model=quantized_model,
+                                                                                      dataset=validation_dataset,
+                                                                                      custom_similarity_metrics=xquant_config.custom_similarity_metrics,
+                                                                                      is_validation=True)
+    similarity_metrics = {
+        OUTPUT_SIMILARITY_METRICS_REPR: repr_similarity[0],
+        OUTPUT_SIMILARITY_METRICS_VAL: val_similarity[0],
+        INTERMEDIATE_SIMILARITY_METRICS_REPR: repr_similarity[1],
+        INTERMEDIATE_SIMILARITY_METRICS_VAL: val_similarity[1]
+    }
+
+    # Add a graph of the quantized model with the similarity metrics to TensorBoard for visualization.
+    fw_report_utils.tb_utils.add_graph_to_tensorboard(quantized_model,
+                                                      similarity_metrics,
+                                                      repr_dataset)
+
+    # Save data to a json file.
+    fw_report_utils.dump_report_to_json(report_dir=xquant_config.report_dir,
+                                        collected_data=similarity_metrics)
+
+    return similarity_metrics

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/dataset_utils.py

@@ -0,0 +1,43 @@
+#  Copyright 2024 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.
+#  ==============================================================================
+
+from typing import Callable
+
+from model_compression_toolkit.logger import Logger
+
+
+class DatasetUtils:
+    """
+    Class with helpful methods for handling different kinds of datasets from the user.
+    """
+
+    @staticmethod
+    def prepare_dataset(dataset: Callable, is_validation: bool, device: str = None):
+        """
+        Prepare the dataset so calling it will return only inputs for the model (like in the case
+        of the representative dataset). For example, when the validation dataset is used, the labels
+        should be removed.
+
+        Args:
+            dataset: Dataset to prepare.
+            is_validation: Whether it's validation dataset or not.
+            device: Device to transfer the data to.
+
+        Returns:
+            Generator to use for retrieving the dataset inputs.
+        """
+
+        Logger.critical("This method should be implemented by the framework-specific DatasetUtils.")  # pragma: no cover
+

mct-nightly-2.1.0.20240618.432/model_compression_toolkit/xquant/common/framework_report_utils.py

@@ -0,0 +1,89 @@
+#  Copyright 2024 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 json
+import os
+
+from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from typing import Any, Dict
+
+from model_compression_toolkit.xquant.common.constants import REPORT_FILENAME
+from model_compression_toolkit.xquant.common.dataset_utils import DatasetUtils
+from model_compression_toolkit.xquant.common.model_folding_utils import ModelFoldingUtils
+from model_compression_toolkit.xquant.common.similarity_calculator import SimilarityCalculator
+from model_compression_toolkit.xquant.common.tensorboard_utils import TensorboardUtils
+from model_compression_toolkit.logger import Logger
+
+
+class FrameworkReportUtils:
+    """
+    Class with various utility components required for generating the report in a specific framework.
+    """
+
+    def __init__(self,
+                 fw_info: FrameworkInfo,
+                 fw_impl: FrameworkImplementation,
+                 similarity_calculator: SimilarityCalculator,
+                 dataset_utils: DatasetUtils,
+                 model_folding_utils: ModelFoldingUtils,
+                 tb_utils: TensorboardUtils):
+        """
+        Initializes the FrameworkReportUtils class with various utility components required for generating the report.
+
+        Args:
+            fw_info (FrameworkInfo): Information about the framework being used.
+            fw_impl (FrameworkImplementation): The implemented functions of the framework.
+            similarity_calculator (SimilarityCalculator): A utility for calculating similarity metrics.
+            dataset_utils (DatasetUtils): Utilities for handling datasets.
+            model_folding_utils (ModelFoldingUtils): Utilities for model folding operations.
+            tb_utils (TensorboardUtils): Utilities for TensorBoard operations.
+        """
+        self.fw_info = fw_info
+        self.fw_impl = fw_impl
+        self.similarity_calculator = similarity_calculator
+        self.dataset_utils = dataset_utils
+        self.model_folding_utils = model_folding_utils
+        self.tb_utils = tb_utils
+
+    def create_report_directory(self, dir_path: str):
+        """
+        Create a directory for saving reports.
+
+        Args:
+            dir_path (str): The path to the directory to create.
+
+        """
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path, exist_ok=True)
+            Logger.info(f"Directory created at: {dir_path}")
+
+    def dump_report_to_json(self,
+                            report_dir: str,
+                            collected_data: Dict[str, Any]):
+        """
+        Dump the collected data (similarity, etc.) into a JSON file.
+
+        Args:
+            report_dir (str): Directory where the report will be saved.
+            collected_data (Dict[str, Any]): Data collected during report generation.
+
+        """
+        report_file_name = os.path.join(report_dir, REPORT_FILENAME)
+        report_file_name = os.path.abspath(report_file_name)
+        Logger.info(f"Dumping report data to: {report_file_name}")
+
+        with open(report_file_name, 'w') as f:
+            json.dump(collected_data, f, indent=4)