mct-nightly 2.2.0.20240922.531__py3-none-any.whl → 2.2.0.20240924.500__py3-none-any.whl

{mct_nightly-2.2.0.20240922.531.dist-info → mct_nightly-2.2.0.20240924.500.dist-info}/METADATA

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

{mct_nightly-2.2.0.20240922.531.dist-info → mct_nightly-2.2.0.20240924.500.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=_e4Qr6XzxLK1myPixyImGS9eeKdlPUijnEQfDZa3iRY,1573
+model_compression_toolkit/__init__.py,sha256=b0vKCDAtdImquoTxXFT-DqKFsjnvUU3tVHOzyTMgiAI,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
@@ -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=Oyro2qg7Bz8TFoimHtrn3JCwHEO9iCrTMy4HktaYZzg,18937
+model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py,sha256=QfMulTLC6X_0Iwvk_VChFIoSdeoiEeJ_rf2IQi5TjBk,19353
 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
@@ -325,7 +325,7 @@ model_compression_toolkit/exporter/model_exporter/keras/mctq_keras_exporter.py,s
 model_compression_toolkit/exporter/model_exporter/pytorch/__init__.py,sha256=uZ2RigbY9O2PJ0Il8wPpS_s7frgg9WUGd_SHeKGyl1A,699
 model_compression_toolkit/exporter/model_exporter/pytorch/base_pytorch_exporter.py,sha256=UPVkEUQCMZ4Lld6CRnEOPEmlfe5vcQZG0Q3FwRBodD4,4021
 model_compression_toolkit/exporter/model_exporter/pytorch/export_serialization_format.py,sha256=bPevy6OBqng41PqytBR55e6cBEuyrUS0H8dWX4zgjQ4,967
-model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py,sha256=647EM7a_at-bGB81ES_6qBjValMnRF5-1TP7Uz6kWQg,6692
+model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py,sha256=07yCCypQaQ3FfhrxKhWwa03bnEgJ8lIknafqeYfvzCQ,6803
 model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_torchscript_pytorch_exporter.py,sha256=ksWV2A-Njo-wAxQ_Ye2sLIZXBWJ_WNyjT7-qFFwvV2o,2897
 model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py,sha256=8vYGKa58BkasvoHejYaPwubOJPcW0s-RY79_Kkw0Hy8,6236
 model_compression_toolkit/exporter/model_wrapper/__init__.py,sha256=7CF2zvpTrIEm8qnbuHnLZyTZkwBBxV24V8QA0oxGbh0,1187
@@ -339,7 +339,7 @@ model_compression_toolkit/exporter/model_wrapper/keras/builder/node_to_quantizer
 model_compression_toolkit/exporter/model_wrapper/pytorch/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/exporter/model_wrapper/pytorch/validate_layer.py,sha256=vQUGbCi8_pGoN8DwQ0IblSeN6L9t6Cr0reZNuCbBpkM,3469
 model_compression_toolkit/exporter/model_wrapper/pytorch/builder/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py,sha256=-4AmWgTG9p8sH2mpns-PaRmvM6J853mrhNc0wt-9ovs,6076
+model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py,sha256=0sx6PLcnJ42LHKn79Qx1FOH615YBqM9OJMF6S1W6plE,6255
 model_compression_toolkit/exporter/model_wrapper/pytorch/builder/node_to_quantizer.py,sha256=4sN5z-6BXrTE5Dp2FX_jKO9ty5iZ2r4RM7XvXtDVLSI,9348
 model_compression_toolkit/gptq/__init__.py,sha256=pEgkJvmf05KSw70iLDTz_6LI_2Oi5L8sTN0JsEUpnpk,1445
 model_compression_toolkit/gptq/runner.py,sha256=La12JTYjWyJW0YW4Al4TP1_Xi4JWBCEKw6FR_JQsxe0,5982
@@ -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=1XWIlnmrmEkKgEsEXlSzQOb8PEebC3O-AYHC9Z7EZ6k,14441
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_keras.py,sha256=8wUlc-ZOHsKwh2hjsdofBIX8Q-iUuNtZd1dF2kZCr2c,7106
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_pytorch.py,sha256=ZaFBnIHLuV1fpOcselo0AWu9NUPGfCuEbuvJ2wFtSHE,6125
+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/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/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.20240922.531.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20240922.531.dist-info/METADATA,sha256=szOhv_jXZsOv5E9Q2G8UK4jWeA5haiQtFjvLPgD9KEg,20813
-mct_nightly-2.2.0.20240922.531.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-mct_nightly-2.2.0.20240922.531.dist-info/top_level.txt,sha256=csdfSXhtRnpWYRzjZ-dRLIhOmM2TEdVXUxG05A5fgb8,39
-mct_nightly-2.2.0.20240922.531.dist-info/RECORD,,
+mct_nightly-2.2.0.20240924.500.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20240924.500.dist-info/METADATA,sha256=tjeXI0teyqJv-zewp1zsb7mL8Jc91rIEzyE83F3SzP8,20813
+mct_nightly-2.2.0.20240924.500.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+mct_nightly-2.2.0.20240924.500.dist-info/top_level.txt,sha256=csdfSXhtRnpWYRzjZ-dRLIhOmM2TEdVXUxG05A5fgb8,39
+mct_nightly-2.2.0.20240924.500.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.20240922.000531"
+__version__ = "2.2.0.20240924.000500"

model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -145,9 +145,14 @@ def _run_operation(n: BaseNode,
         else:
             out_tensors_of_n_float = op_func(input_tensors, *op_call_args, **functional_kwargs)
     else:
-        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)
+        if isinstance(op_func, PytorchQuantizationWrapper) and isinstance(n, FunctionalNode) and n.functional_op is not torch.gather:
+            # in wrapped nodes, the op args & kwargs are already in the PytorchQuantizationWrapper.
+            # Temporary patch: for torch.gather this is not the case, so need to merge inputs.
+            out_tensors_of_n_float = op_func(*input_tensors)
+        else:
+            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.
     out_tensors_of_n = out_tensors_of_n_float

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -89,12 +89,12 @@ if FOUND_ONNX:
             else:
                 Logger.info(f"Exporting fake-quant onnx model: {self.save_model_path}")
 
-            model_input = to_torch_tensor(next(self.repr_dataset())[0])
+            model_input = to_torch_tensor(next(self.repr_dataset()))
 
             if hasattr(self.model, 'metadata'):
                 onnx_bytes = BytesIO()
                 torch.onnx.export(self.model,
-                                  model_input,
+                                  tuple(model_input) if isinstance(model_input, list) else model_input,
                                   onnx_bytes,
                                   opset_version=self._onnx_opset_version,
                                   verbose=False,
@@ -107,7 +107,7 @@ if FOUND_ONNX:
                 onnx.save_model(onnx_model, self.save_model_path)
             else:
                 torch.onnx.export(self.model,
-                                  model_input,
+                                  tuple(model_input) if isinstance(model_input, list) else model_input,
                                   self.save_model_path,
                                   opset_version=self._onnx_opset_version,
                                   verbose=False,

model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py

@@ -50,9 +50,11 @@ if FOUND_TORCH:
                               for attr in weight_quantizers if isinstance(attr, int)}
             # When wrapping functional nodes, need to set call args\kwargs in wrapper, because they
             # are used during wrapper call method.
+            # Temporary patch: for torch.gather this is not the case, so args & kwargs shouldn't be
+            # saved in the warpper.
             func_node_kwargs = {OP_CALL_ARGS: node.op_call_args,
                                 OP_CALL_KWARGS: node.op_call_kwargs
-                                } if isinstance(node, FunctionalNode) else {}
+                                } if isinstance(node, FunctionalNode) and not node.functional_op is torch.gather else {}
             return PytorchQuantizationWrapper(module, weight_quantizers, weights_values,
                                               is_inputs_as_list=node.inputs_as_list,
                                               **func_node_kwargs)

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

@@ -24,6 +24,23 @@ from model_compression_toolkit.target_platform_capabilities.target_platform.op_q
 
 tp = mct.target_platform
 
+OPSET_NO_QUANTIZATION = "NoQuantization"
+OPSET_QUANTIZATION_PRESERVING = "QuantizationPreserving"
+OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS = "DimensionManipulationOpsWithWeights"
+OPSET_DIMENSION_MANIPULATION_OPS = "DimensionManipulationOps"
+OPSET_MERGE_OPS = "MergeOps"
+OPSET_CONV = "Conv"
+OPSET_FULLY_CONNECTED = "FullyConnected"
+OPSET_ANY_RELU = "AnyReLU"
+OPSET_ADD = "Add"
+OPSET_SUB = "Sub"
+OPSET_MUL = "Mul"
+OPSET_DIV = "Div"
+OPSET_PRELU = "PReLU"
+OPSET_SWISH = "Swish"
+OPSET_SIGMOID = "Sigmoid"
+OPSET_TANH = "Tanh"
+
 
 def get_tp_model() -> TargetPlatformModel:
     """
@@ -189,6 +206,10 @@ def generate_tp_model(default_config: OpQuantizationConfig,
                                                                                    const_config_input16_per_tensor],
                                                                                   base_config=const_config_input16_per_tensor)
 
+    qpreserving_const_config = const_config.clone_and_edit(enable_activation_quantization=False,
+                                                           quantization_preserving=True)
+    qpreserving_const_config_options = tp.QuantizationConfigOptions([qpreserving_const_config])
+
     # Create a TargetPlatformModel and set its default quantization config.
     # This default configuration will be used for all operations
     # unless specified otherwise (see OperatorsSet, for example):
@@ -207,39 +228,40 @@ def generate_tp_model(default_config: OpQuantizationConfig,
 
         # May suit for operations like: Dropout, Reshape, etc.
         default_qco = tp.get_default_quantization_config_options()
-        tp.OperatorsSet("NoQuantization",
+        tp.OperatorsSet(OPSET_NO_QUANTIZATION,
                         default_qco.clone_and_edit(enable_activation_quantization=False)
                         .clone_and_edit_weight_attribute(enable_weights_quantization=False))
-        tp.OperatorsSet("QuantizationPreserving",
+        tp.OperatorsSet(OPSET_QUANTIZATION_PRESERVING,
                         default_qco.clone_and_edit(enable_activation_quantization=False,
                                                    quantization_preserving=True)
                         .clone_and_edit_weight_attribute(enable_weights_quantization=False))
-        tp.OperatorsSet("DimensionManipulationOps",
+        tp.OperatorsSet(OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, qpreserving_const_config_options)
+        tp.OperatorsSet(OPSET_DIMENSION_MANIPULATION_OPS,
                         default_qco.clone_and_edit(enable_activation_quantization=False,
                                                    quantization_preserving=True,
                                                    supported_input_activation_n_bits=(8, 16))
                         .clone_and_edit_weight_attribute(enable_weights_quantization=False))
-        tp.OperatorsSet("MergeOps", const_configuration_options_inout16_per_tensor)
+        tp.OperatorsSet(OPSET_MERGE_OPS, const_configuration_options_inout16_per_tensor)
 
         # Create Mixed-Precision quantization configuration options from the given list of OpQuantizationConfig objects
         mixed_precision_configuration_options = tp.QuantizationConfigOptions(mixed_precision_cfg_list,
                                                                              base_config=base_config)
 
         # Define operator sets that use mixed_precision_configuration_options:
-        conv = tp.OperatorsSet("Conv", mixed_precision_configuration_options)
-        fc = tp.OperatorsSet("FullyConnected", mixed_precision_configuration_options)
+        conv = tp.OperatorsSet(OPSET_CONV, mixed_precision_configuration_options)
+        fc = tp.OperatorsSet(OPSET_FULLY_CONNECTED, mixed_precision_configuration_options)
 
         # Define operations sets without quantization configuration
         # options (useful for creating fusing patterns, for example):
-        any_relu = tp.OperatorsSet("AnyReLU")
-        add = tp.OperatorsSet("Add", const_configuration_options_inout16)
-        sub = tp.OperatorsSet("Sub", const_configuration_options_inout16)
-        mul = tp.OperatorsSet("Mul", const_configuration_options_inout16)
-        div = tp.OperatorsSet("Div", const_configuration_options)
-        prelu = tp.OperatorsSet("PReLU")
-        swish = tp.OperatorsSet("Swish")
-        sigmoid = tp.OperatorsSet("Sigmoid")
-        tanh = tp.OperatorsSet("Tanh")
+        any_relu = tp.OperatorsSet(OPSET_ANY_RELU)
+        add = tp.OperatorsSet(OPSET_ADD, const_configuration_options_inout16)
+        sub = tp.OperatorsSet(OPSET_SUB, const_configuration_options_inout16)
+        mul = tp.OperatorsSet(OPSET_MUL, const_configuration_options_inout16)
+        div = tp.OperatorsSet(OPSET_DIV, const_configuration_options)
+        prelu = tp.OperatorsSet(OPSET_PRELU)
+        swish = tp.OperatorsSet(OPSET_SWISH)
+        sigmoid = tp.OperatorsSet(OPSET_SIGMOID)
+        tanh = tp.OperatorsSet(OPSET_TANH)
 
         # Combine multiple operators into a single operator to avoid quantization between
         # them. To do this we define fusing patterns using the OperatorsSets that were created.

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

@@ -35,6 +35,10 @@ else:
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tp_model import get_tp_model
 import model_compression_toolkit as mct
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4 import __version__ as TPC_VERSION
+from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tp_model import OPSET_NO_QUANTIZATION, \
+    OPSET_QUANTIZATION_PRESERVING, OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, OPSET_DIMENSION_MANIPULATION_OPS, \
+    OPSET_MERGE_OPS, OPSET_CONV, OPSET_FULLY_CONNECTED, OPSET_ANY_RELU, OPSET_ADD, OPSET_SUB, OPSET_MUL, OPSET_DIV, \
+    OPSET_PRELU, OPSET_SWISH, OPSET_SIGMOID, OPSET_TANH
 
 tp = mct.target_platform
 
@@ -74,10 +78,8 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                Dropout,
                                MaxPooling2D,
                                tf.split,
-                               tf.gather,
                                tf.cast,
                                tf.unstack,
-                               tf.compat.v1.gather,
                                tf.__operators__.getitem,
                                tf.strided_slice]
     quantization_preserving_list_16bit_input = [Reshape,
@@ -90,11 +92,12 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
         no_quant_list.append(SSDPostProcess)
 
     with keras_tpc:
-        tp.OperationsSetToLayers("NoQuantization", no_quant_list)
-        tp.OperationsSetToLayers("QuantizationPreserving", quantization_preserving)
-        tp.OperationsSetToLayers("DimensionManipulationOps", quantization_preserving_list_16bit_input)
-        tp.OperationsSetToLayers("MergeOps", [tf.stack, tf.concat, Concatenate])
-        tp.OperationsSetToLayers("Conv",
+        tp.OperationsSetToLayers(OPSET_NO_QUANTIZATION, no_quant_list)
+        tp.OperationsSetToLayers(OPSET_QUANTIZATION_PRESERVING, quantization_preserving)
+        tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS, quantization_preserving_list_16bit_input)
+        tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, [tf.gather, tf.compat.v1.gather])
+        tp.OperationsSetToLayers(OPSET_MERGE_OPS, [tf.stack, tf.concat, Concatenate])
+        tp.OperationsSetToLayers(OPSET_CONV,
                                  [Conv2D,
                                   DepthwiseConv2D,
                                   Conv2DTranspose,
@@ -111,23 +114,23 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                          DepthwiseConv2D: KERAS_DEPTHWISE_KERNEL,
                                          tf.nn.depthwise_conv2d: KERAS_DEPTHWISE_KERNEL}, default_value=KERAS_KERNEL),
                                      BIAS_ATTR: DefaultDict(default_value=BIAS)})
-        tp.OperationsSetToLayers("FullyConnected", [Dense],
+        tp.OperationsSetToLayers(OPSET_FULLY_CONNECTED, [Dense],
                                  attr_mapping={KERNEL_ATTR: DefaultDict(default_value=KERAS_KERNEL),
                                                BIAS_ATTR: DefaultDict(default_value=BIAS)})
-        tp.OperationsSetToLayers("AnyReLU", [tf.nn.relu,
-                                             tf.nn.relu6,
-                                             tf.nn.leaky_relu,
-                                             ReLU,
-                                             LeakyReLU,
-                                             tp.LayerFilterParams(Activation, activation="relu"),
-                                             tp.LayerFilterParams(Activation, activation="leaky_relu")])
-        tp.OperationsSetToLayers("Add", [tf.add, Add])
-        tp.OperationsSetToLayers("Sub", [tf.subtract, Subtract])
-        tp.OperationsSetToLayers("Mul", [tf.math.multiply, Multiply])
-        tp.OperationsSetToLayers("Div", [tf.math.divide, tf.math.truediv])
-        tp.OperationsSetToLayers("PReLU", [PReLU])
-        tp.OperationsSetToLayers("Swish", [tf.nn.swish, tp.LayerFilterParams(Activation, activation="swish")])
-        tp.OperationsSetToLayers("Sigmoid", [tf.nn.sigmoid, tp.LayerFilterParams(Activation, activation="sigmoid")])
-        tp.OperationsSetToLayers("Tanh", [tf.nn.tanh, tp.LayerFilterParams(Activation, activation="tanh")])
+        tp.OperationsSetToLayers(OPSET_ANY_RELU, [tf.nn.relu,
+                                                  tf.nn.relu6,
+                                                  tf.nn.leaky_relu,
+                                                  ReLU,
+                                                  LeakyReLU,
+                                                  tp.LayerFilterParams(Activation, activation="relu"),
+                                                  tp.LayerFilterParams(Activation, activation="leaky_relu")])
+        tp.OperationsSetToLayers(OPSET_ADD, [tf.add, Add])
+        tp.OperationsSetToLayers(OPSET_SUB, [tf.subtract, Subtract])
+        tp.OperationsSetToLayers(OPSET_MUL, [tf.math.multiply, Multiply])
+        tp.OperationsSetToLayers(OPSET_DIV, [tf.math.divide, tf.math.truediv])
+        tp.OperationsSetToLayers(OPSET_PRELU, [PReLU])
+        tp.OperationsSetToLayers(OPSET_SWISH, [tf.nn.swish, tp.LayerFilterParams(Activation, activation="swish")])
+        tp.OperationsSetToLayers(OPSET_SIGMOID, [tf.nn.sigmoid, tp.LayerFilterParams(Activation, activation="sigmoid")])
+        tp.OperationsSetToLayers(OPSET_TANH, [tf.nn.tanh, tp.LayerFilterParams(Activation, activation="tanh")])
 
     return keras_tpc

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

@@ -29,6 +29,10 @@ from model_compression_toolkit.target_platform_capabilities.constants import KER
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tp_model import get_tp_model
 import model_compression_toolkit as mct
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4 import __version__ as TPC_VERSION
+from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tp_model import OPSET_NO_QUANTIZATION, \
+    OPSET_QUANTIZATION_PRESERVING, OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, OPSET_DIMENSION_MANIPULATION_OPS, \
+    OPSET_MERGE_OPS, OPSET_CONV, OPSET_FULLY_CONNECTED, OPSET_ANY_RELU, OPSET_ADD, OPSET_SUB, OPSET_MUL, OPSET_DIV, \
+    OPSET_PRELU, OPSET_SWISH, OPSET_SIGMOID, OPSET_TANH
 
 tp = mct.target_platform
 
@@ -65,49 +69,49 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                    BIAS_ATTR: DefaultDict(default_value=BIAS)}
 
     with pytorch_tpc:
-        tp.OperationsSetToLayers("NoQuantization", [torch.Tensor.size,
-                                                    equal,
-                                                    argmax,
-                                                    topk])
-        tp.OperationsSetToLayers("QuantizationPreserving", [Dropout,
-                                                            dropout,
-                                                            split,
-                                                            chunk,
-                                                            unbind,
-                                                            gather,
-                                                            MaxPool2d])
-        tp.OperationsSetToLayers("DimensionManipulationOps", [Flatten,
-                                                                      flatten,
-                                                                      operator.getitem,
-                                                                      reshape,
-                                                                      unsqueeze,
-                                                                      squeeze,
-                                                                      permute,
-                                                                      transpose])
-        tp.OperationsSetToLayers("MergeOps",
+        tp.OperationsSetToLayers(OPSET_NO_QUANTIZATION, [torch.Tensor.size,
+                                                         equal,
+                                                         argmax,
+                                                         topk])
+        tp.OperationsSetToLayers(OPSET_QUANTIZATION_PRESERVING, [Dropout,
+                                                                 dropout,
+                                                                 split,
+                                                                 chunk,
+                                                                 unbind,
+                                                                 MaxPool2d])
+        tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS, [Flatten,
+                                                                    flatten,
+                                                                    operator.getitem,
+                                                                    reshape,
+                                                                    unsqueeze,
+                                                                    squeeze,
+                                                                    permute,
+                                                                    transpose])
+        tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, [gather])
+        tp.OperationsSetToLayers(OPSET_MERGE_OPS,
                                  [torch.stack, torch.cat, torch.concat, torch.concatenate])
 
-        tp.OperationsSetToLayers("Conv", [Conv2d, ConvTranspose2d],
+        tp.OperationsSetToLayers(OPSET_CONV, [Conv2d, ConvTranspose2d],
                                  attr_mapping=pytorch_linear_attr_mapping)
-        tp.OperationsSetToLayers("FullyConnected", [Linear],
+        tp.OperationsSetToLayers(OPSET_FULLY_CONNECTED, [Linear],
                                  attr_mapping=pytorch_linear_attr_mapping)
-        tp.OperationsSetToLayers("AnyReLU", [torch.relu,
-                                             ReLU,
-                                             ReLU6,
-                                             LeakyReLU,
-                                             relu,
-                                             relu6,
-                                             leaky_relu,
-                                             tp.LayerFilterParams(Hardtanh, min_val=0),
-                                             tp.LayerFilterParams(hardtanh, min_val=0)])
+        tp.OperationsSetToLayers(OPSET_ANY_RELU, [torch.relu,
+                                                  ReLU,
+                                                  ReLU6,
+                                                  LeakyReLU,
+                                                  relu,
+                                                  relu6,
+                                                  leaky_relu,
+                                                  tp.LayerFilterParams(Hardtanh, min_val=0),
+                                                  tp.LayerFilterParams(hardtanh, min_val=0)])
 
-        tp.OperationsSetToLayers("Add", [operator.add, add])
-        tp.OperationsSetToLayers("Sub", [operator.sub, sub, subtract])
-        tp.OperationsSetToLayers("Mul", [operator.mul, mul, multiply])
-        tp.OperationsSetToLayers("Div", [operator.truediv, div, divide])
-        tp.OperationsSetToLayers("PReLU", [PReLU, prelu])
-        tp.OperationsSetToLayers("Swish", [SiLU, silu, Hardswish, hardswish])
-        tp.OperationsSetToLayers("Sigmoid", [Sigmoid, sigmoid])
-        tp.OperationsSetToLayers("Tanh", [Tanh, tanh])
+        tp.OperationsSetToLayers(OPSET_ADD, [operator.add, add])
+        tp.OperationsSetToLayers(OPSET_SUB, [operator.sub, sub, subtract])
+        tp.OperationsSetToLayers(OPSET_MUL, [operator.mul, mul, multiply])
+        tp.OperationsSetToLayers(OPSET_DIV, [operator.truediv, div, divide])
+        tp.OperationsSetToLayers(OPSET_PRELU, [PReLU, prelu])
+        tp.OperationsSetToLayers(OPSET_SWISH, [SiLU, silu, Hardswish, hardswish])
+        tp.OperationsSetToLayers(OPSET_SIGMOID, [Sigmoid, sigmoid])
+        tp.OperationsSetToLayers(OPSET_TANH, [Tanh, tanh])
 
     return pytorch_tpc