mct-nightly 2.2.0.20250106.546__py3-none-any.whl → 2.2.0.20250107.15510__py3-none-any.whl

model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/attach2pytorch.py

@@ -18,74 +18,89 @@ import operator
 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, minimum, \
-    maximum
-from torch.nn import Conv2d, Linear, ConvTranspose2d, MaxPool2d, BatchNorm2d
-from torch.nn import Dropout, Flatten, Hardtanh
-from torch.nn import ReLU, ReLU6, PReLU, SiLU, Sigmoid, Tanh, Hardswish, Hardsigmoid, LeakyReLU, GELU
+    maximum, softmax, fake_quantize_per_channel_affine
+from torch.nn import Conv2d, Linear, ConvTranspose2d, MaxPool2d, BatchNorm2d, Dropout, Flatten, Hardtanh, ReLU, ReLU6, \
+    PReLU, SiLU, Sigmoid, Tanh, Hardswish, Hardsigmoid, LeakyReLU, GELU, LogSoftmax, Softmax, ELU, AvgPool2d, ZeroPad2d
+from torch.nn.functional import relu, relu6, prelu, silu, hardtanh, hardswish, hardsigmoid, leaky_relu, gelu, fold
 import torch.nn.functional as F
-from torch.nn.functional import relu, relu6, prelu, silu, hardtanh, hardswish, hardsigmoid, leaky_relu, gelu
 
 from model_compression_toolkit import DefaultDict
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, PYTORCH_KERNEL, BIAS, \
     BIAS_ATTR
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import OperatorSetNames
-from model_compression_toolkit.target_platform_capabilities.target_platform import LayerFilterParams
+from model_compression_toolkit.target_platform_capabilities.target_platform import LayerFilterParams, Eq
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework.attach2fw import \
-    AttachTpModelToFw
+    AttachTpcToFramework
 
 
-class AttachTpModelToPytorch(AttachTpModelToFw):
+class AttachTpcToPytorch(AttachTpcToFramework):
     def __init__(self):
         super().__init__()
 
         self._opset2layer = {
-            OperatorSetNames.OPSET_CONV.value: [Conv2d],
-            OperatorSetNames.OPSET_CONV_TRANSPOSE.value: [ConvTranspose2d],
-            OperatorSetNames.OPSET_FULLY_CONNECTED.value: [Linear],
-            OperatorSetNames.OPSET_CONCATENATE.value: [torch.cat, torch.concat, torch.concatenate],
-            OperatorSetNames.OPSET_STACK.value: [torch.stack],
-            OperatorSetNames.OPSET_UNSTACK.value: [unbind],
-            OperatorSetNames.OPSET_GATHER.value: [gather],
-            OperatorSetNames.OPSET_EXPAND.value: [torch.Tensor.expand],
-            OperatorSetNames.OPSET_BATCH_NORM.value: [BatchNorm2d],
-            OperatorSetNames.OPSET_RELU.value: [torch.relu, ReLU, relu],
-            OperatorSetNames.OPSET_RELU6.value: [ReLU6, relu6],
-            OperatorSetNames.OPSET_LEAKY_RELU.value: [LeakyReLU, leaky_relu],
-            OperatorSetNames.OPSET_HARD_TANH.value: [LayerFilterParams(Hardtanh, min_val=0),
-                                                     LayerFilterParams(hardtanh, min_val=0)],
-            OperatorSetNames.OPSET_ADD.value: [operator.add, add],
-            OperatorSetNames.OPSET_SUB.value: [operator.sub, sub, subtract],
-            OperatorSetNames.OPSET_MUL.value: [operator.mul, mul, multiply],
-            OperatorSetNames.OPSET_DIV.value: [operator.truediv, div, divide],
-            OperatorSetNames.OPSET_MIN.value: [minimum],
-            OperatorSetNames.OPSET_MAX.value: [maximum],
-            OperatorSetNames.OPSET_PRELU.value: [PReLU, prelu],
-            OperatorSetNames.OPSET_SWISH.value: [SiLU, silu],
-            OperatorSetNames.OPSET_SIGMOID.value: [Sigmoid, sigmoid, F.sigmoid],
-            OperatorSetNames.OPSET_TANH.value: [Tanh, tanh, F.tanh],
-            OperatorSetNames.OPSET_GELU.value: [GELU, gelu],
-            OperatorSetNames.OPSET_HARDSIGMOID.value: [Hardsigmoid, hardsigmoid],
-            OperatorSetNames.OPSET_HARDSWISH.value: [Hardswish, hardswish],
-            OperatorSetNames.OPSET_FLATTEN.value: [Flatten, flatten],
-            OperatorSetNames.OPSET_GET_ITEM.value: [operator.getitem],
-            OperatorSetNames.OPSET_RESHAPE.value: [reshape],
-            OperatorSetNames.OPSET_UNSQUEEZE.value: [unsqueeze],
-            OperatorSetNames.OPSET_SQUEEZE.value: [squeeze],
-            OperatorSetNames.OPSET_PERMUTE.value: [permute],
-            OperatorSetNames.OPSET_TRANSPOSE.value: [transpose],
-            OperatorSetNames.OPSET_DROPOUT.value: [Dropout, dropout],
-            OperatorSetNames.OPSET_SPLIT.value: [split],
-            OperatorSetNames.OPSET_CHUNK.value: [chunk],
-            OperatorSetNames.OPSET_MAXPOOL.value: [MaxPool2d],
-            OperatorSetNames.OPSET_SIZE.value: [torch.Tensor.size],
-            OperatorSetNames.OPSET_SHAPE.value: [torch.Tensor.shape],
-            OperatorSetNames.OPSET_EQUAL.value: [equal],
-            OperatorSetNames.OPSET_ARGMAX.value: [argmax],
-            OperatorSetNames.OPSET_TOPK.value: [topk],
+            OperatorSetNames.CONV: [Conv2d],
+            OperatorSetNames.DEPTHWISE_CONV: [],  # no specific operator for depthwise conv in pytorch
+            OperatorSetNames.CONV_TRANSPOSE: [ConvTranspose2d],
+            OperatorSetNames.FULLY_CONNECTED: [Linear],
+            OperatorSetNames.CONCATENATE: [torch.cat, torch.concat, torch.concatenate],
+            OperatorSetNames.STACK: [torch.stack],
+            OperatorSetNames.UNSTACK: [unbind],
+            OperatorSetNames.GATHER: [gather],
+            OperatorSetNames.EXPAND: [torch.Tensor.expand],
+            OperatorSetNames.BATCH_NORM: [BatchNorm2d],
+            OperatorSetNames.RELU: [torch.relu, ReLU, relu],
+            OperatorSetNames.RELU6: [ReLU6, relu6],
+            OperatorSetNames.LEAKY_RELU: [LeakyReLU, leaky_relu],
+            OperatorSetNames.HARD_TANH: [LayerFilterParams(Hardtanh, min_val=0),
+                                         LayerFilterParams(hardtanh, min_val=0)],
+            OperatorSetNames.ADD: [operator.add, add],
+            OperatorSetNames.SUB: [operator.sub, sub, subtract],
+            OperatorSetNames.MUL: [operator.mul, mul, multiply],
+            OperatorSetNames.DIV: [operator.truediv, div, divide],
+            OperatorSetNames.ADD_BIAS: [],  # no specific operator for bias_add in pytorch
+            OperatorSetNames.MIN: [minimum],
+            OperatorSetNames.MAX: [maximum],
+            OperatorSetNames.PRELU: [PReLU, prelu],
+            OperatorSetNames.SWISH: [SiLU, silu],
+            OperatorSetNames.SIGMOID: [Sigmoid, sigmoid, F.sigmoid],
+            OperatorSetNames.TANH: [Tanh, tanh, F.tanh],
+            OperatorSetNames.GELU: [GELU, gelu],
+            OperatorSetNames.HARDSIGMOID: [Hardsigmoid, hardsigmoid],
+            OperatorSetNames.HARDSWISH: [Hardswish, hardswish],
+            OperatorSetNames.FLATTEN: [Flatten, flatten],
+            OperatorSetNames.GET_ITEM: [operator.getitem],
+            OperatorSetNames.RESHAPE: [reshape],
+            OperatorSetNames.UNSQUEEZE: [unsqueeze],
+            OperatorSetNames.SQUEEZE: [squeeze],
+            OperatorSetNames.PERMUTE: [permute],
+            OperatorSetNames.TRANSPOSE: [transpose],
+            OperatorSetNames.DROPOUT: [Dropout, dropout],
+            OperatorSetNames.SPLIT_CHUNK: [split, chunk],
+            OperatorSetNames.MAXPOOL: [MaxPool2d, F.max_pool2d],
+            OperatorSetNames.AVGPOOL: [AvgPool2d, F.avg_pool2d],
+            OperatorSetNames.SIZE: [torch.Tensor.size],
+            OperatorSetNames.RESIZE: [torch.Tensor.resize],
+            OperatorSetNames.PAD: [F.pad],
+            OperatorSetNames.FOLD: [fold],
+            OperatorSetNames.SHAPE: [torch.Tensor.shape],
+            OperatorSetNames.EQUAL: [equal],
+            OperatorSetNames.ARGMAX: [argmax],
+            OperatorSetNames.TOPK: [topk],
+            OperatorSetNames.FAKE_QUANT: [fake_quantize_per_channel_affine],
+            OperatorSetNames.ZERO_PADDING2D: [ZeroPad2d],
+            OperatorSetNames.CAST: [torch.Tensor.type],
+            OperatorSetNames.STRIDED_SLICE: [],  # no such operator in pytorch, the equivalent is get_item which has a separate operator set
+            OperatorSetNames.ELU: [ELU, F.elu],
+            OperatorSetNames.SOFTMAX: [Softmax, softmax, F.softmax],
+            OperatorSetNames.LOG_SOFTMAX: [LogSoftmax],
+            OperatorSetNames.L2NORM: [LayerFilterParams(torch.nn.functional.normalize,
+                                                        Eq('p', 2) | Eq('p', None))],
+            OperatorSetNames.SSD_POST_PROCESS: [],  # no such operator in pytorch
+            OperatorSetNames.COMBINED_NON_MAX_SUPPRESSION: []  # no such operator in pytorch
         }
 
         pytorch_linear_attr_mapping = {KERNEL_ATTR: DefaultDict(default_value=PYTORCH_KERNEL),
                                        BIAS_ATTR: DefaultDict(default_value=BIAS)}
-        self._opset2attr_mapping = {OperatorSetNames.OPSET_CONV.value: pytorch_linear_attr_mapping,
-                                    OperatorSetNames.OPSET_CONV_TRANSPOSE.value: pytorch_linear_attr_mapping,
-                                    OperatorSetNames.OPSET_FULLY_CONNECTED.value: pytorch_linear_attr_mapping}
+        self._opset2attr_mapping = {OperatorSetNames.CONV: pytorch_linear_attr_mapping,
+                                    OperatorSetNames.CONV_TRANSPOSE: pytorch_linear_attr_mapping,
+                                    OperatorSetNames.FULLY_CONNECTED: pytorch_linear_attr_mapping}

model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/operations_to_layers.py

@@ -138,10 +138,8 @@ class OperationsToLayers:
                               OperationsSetToLayers), f'Operators set should be of type OperationsSetToLayers but it ' \
                                                       f'is of type {type(ops2layers)}'
 
-            # Assert that opset in the current TargetPlatformCapabilities and has a unique name.
-            opset_in_model = is_opset_in_model(_current_tpc.get().tp_model, ops2layers.name)
-            assert opset_in_model, f'{ops2layers.name} is not defined in the target platform model that is associated with the target platform capabilities.'
-            assert not (ops2layers.name in existing_opset_names), f'OperationsSetToLayers names should be unique, but {ops2layers.name} appears to violate it.'
+            # Assert that opset has a unique name.
+            assert ops2layers.name not in existing_opset_names, f'OperationsSetToLayers names should be unique, but {ops2layers.name} appears to violate it.'
             existing_opset_names.append(ops2layers.name)
 
             # Assert that a layer does not appear in more than a single OperatorsSet in the TargetPlatformModel.

model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/target_platform_capabilities.py

@@ -156,7 +156,6 @@ class TargetPlatformCapabilities(ImmutableClass):
         if exc_value is not None:
             print(exc_value, exc_value.args)
             raise exc_value
-        self.raise_warnings()
         self.layer2qco, self.filterlayer2qco = self._get_config_options_mapping()
         _current_tpc.reset()
         self.initialized_done()
@@ -226,15 +225,6 @@ class TargetPlatformCapabilities(ImmutableClass):
         if opset_to_remove in self.__tp_model_opsets_not_used:
             self.__tp_model_opsets_not_used.remove(opset_to_remove)
 
-    def raise_warnings(self):
-        """
-
-        Log warnings regards unused opsets.
-
-        """
-        for op in self.__tp_model_opsets_not_used:
-            Logger.warning(f'{op} is defined in TargetPlatformModel, but is not used in TargetPlatformCapabilities.')
-
     @property
     def is_simd_padding(self) -> bool:
         """

model_compression_toolkit/target_platform_capabilities/tpc_io_handler.py

@@ -0,0 +1,93 @@
+# 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 pathlib import Path
+from typing import Union
+
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
+import json
+
+
+def load_target_platform_model(tp_model_or_path: Union[TargetPlatformModel, str]) -> TargetPlatformModel:
+    """
+        Parses the tp_model input, which can be either a TargetPlatformModel object
+        or a string path to a JSON file.
+
+        Parameters:
+            tp_model_or_path (Union[TargetPlatformModel, str]): Input target platform model or path to .JSON file.
+
+        Returns:
+            TargetPlatformModel: The parsed TargetPlatformModel.
+
+        Raises:
+            FileNotFoundError: If the JSON file does not exist.
+            ValueError: If the JSON content is invalid or cannot initialize the TargetPlatformModel.
+            TypeError: If the input is neither a TargetPlatformModel nor a valid JSON file path.
+        """
+    if isinstance(tp_model_or_path, TargetPlatformModel):
+        return tp_model_or_path
+
+    if isinstance(tp_model_or_path, str):
+        path = Path(tp_model_or_path)
+
+        if not path.exists() or not path.is_file():
+            raise FileNotFoundError(f"The path '{tp_model_or_path}' is not a valid file.")
+        # Verify that the file has a .json extension
+        if path.suffix.lower() != '.json':
+            raise ValueError(f"The file '{path}' does not have a '.json' extension.")
+        try:
+            with path.open('r', encoding='utf-8') as file:
+                data = file.read()
+        except OSError as e:
+            raise ValueError(f"Error reading the file '{tp_model_or_path}': {e.strerror}.") from e
+
+        try:
+            return TargetPlatformModel.parse_raw(data)
+        except ValueError as e:
+            raise ValueError(f"Invalid JSON for loading TargetPlatformModel in '{tp_model_or_path}': {e}.") from e
+        except Exception as e:
+            raise ValueError(f"Unexpected error while initializing TargetPlatformModel: {e}.") from e
+
+    raise TypeError(
+        f"tp_model_or_path must be either a TargetPlatformModel instance or a string path to a JSON file, "
+        f"but received type '{type(tp_model_or_path).__name__}'."
+    )
+
+
+def export_target_platform_model(model: TargetPlatformModel, export_path: Union[str, Path]) -> None:
+    """
+    Exports a TargetPlatformModel instance to a JSON file.
+
+    Parameters:
+        model (TargetPlatformModel): The TargetPlatformModel instance to export.
+        export_path (Union[str, Path]): The file path to export the model to.
+
+    Raises:
+        ValueError: If the model is not an instance of TargetPlatformModel.
+        OSError: If there is an issue writing to the file.
+    """
+    if not isinstance(model, TargetPlatformModel):
+        raise ValueError("The provided model is not a valid TargetPlatformModel instance.")
+
+    path = Path(export_path)
+    try:
+        # Ensure the parent directory exists
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Export the model to JSON and write to the file
+        with path.open('w', encoding='utf-8') as file:
+            file.write(model.json(indent=4))
+    except OSError as e:
+        raise OSError(f"Failed to write to file '{export_path}': {e.strerror}") from e

model_compression_toolkit/target_platform_capabilities/tpc_models/get_target_platform_capabilities.py

@@ -12,45 +12,63 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from model_compression_toolkit.constants import TENSORFLOW, PYTORCH
+from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL, IMX500_TP_MODEL, \
+    TFLITE_TP_MODEL, QNNPACK_TP_MODEL
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
 
-from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tp_model import get_tp_model as get_tp_model_imx500_v1
+from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tp_model import get_tp_model as get_tp_model_tflite_v1
+from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tp_model import get_tp_model as get_tp_model_qnnpack_v1
 
-from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.target_platform_capabilities import \
-    get_tpc_dict_by_fw as get_imx500_tpc
-from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.target_platform_capabilities import \
-    get_tpc_dict_by_fw as get_tflite_tpc
-from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.target_platform_capabilities import \
-    get_tpc_dict_by_fw as get_qnnpack_tpc
-from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL, IMX500_TP_MODEL, TFLITE_TP_MODEL, QNNPACK_TP_MODEL,  LATEST
-
-tpc_dict = {DEFAULT_TP_MODEL: get_imx500_tpc,
-            IMX500_TP_MODEL: get_imx500_tpc,
-            TFLITE_TP_MODEL: get_tflite_tpc,
-            QNNPACK_TP_MODEL: get_qnnpack_tpc}
 
+# TODO: These methods need to be replaced once modifying the TPC API.
 
 def get_target_platform_capabilities(fw_name: str,
                                      target_platform_name: str,
-                                     target_platform_version: str = None) -> TargetPlatformCapabilities:
+                                     target_platform_version: str = None) -> TargetPlatformModel:
     """
-    Get a TargetPlatformCapabilities by the target platform model name and the framework name.
-    For now, it supports frameworks 'tensorflow' and 'pytorch'. For both of them
-    the target platform model can be 'default', 'imx500', 'tflite', or 'qnnpack'.
+    This is a degenerated function that only returns the MCT default TargetPlatformModel object, to comply with the
+    existing TPC API.
 
     Args:
         fw_name: Framework name of the TargetPlatformCapabilities.
         target_platform_name: Target platform model name the model will use for inference.
         target_platform_version: Target platform capabilities version.
+
     Returns:
-        A TargetPlatformCapabilities object that models the hardware and attaches
-        a framework information to it.
+        A default TargetPlatformModel object.
     """
-    assert target_platform_name in tpc_dict, f'Target platform {target_platform_name} is not defined!'
-    fw_tpc = tpc_dict.get(target_platform_name)
-    tpc_versions = fw_tpc(fw_name)
-    if target_platform_version is None:
-        target_platform_version = LATEST
-    else:
-        assert target_platform_version in tpc_versions, (f'TPC version {target_platform_version} is not supported for '
-                                                         f'framework {fw_name}.')
-    return tpc_versions[target_platform_version]()
+
+    assert fw_name in [TENSORFLOW, PYTORCH], f"Unsupported framework {fw_name}."
+
+    if target_platform_name == DEFAULT_TP_MODEL:
+        return get_tp_model_imx500_v1()
+
+    assert target_platform_version == 'v1' or target_platform_version is None, \
+        "The usage of get_target_platform_capabilities API is supported only with the default TPC ('v1')."
+
+    if target_platform_name == IMX500_TP_MODEL:
+        return get_tp_model_imx500_v1()
+    elif target_platform_name == TFLITE_TP_MODEL:
+        return get_tp_model_tflite_v1()
+    elif target_platform_name == QNNPACK_TP_MODEL:
+        return get_tp_model_qnnpack_v1()
+
+    raise ValueError(f"Unsupported target platform name {target_platform_name}.")
+
+
+def get_tpc_model(name: str, tp_model: TargetPlatformModel):
+    """
+    This is a utility method that just returns the TargetPlatformModel that it receives, to support existing TPC API.
+
+    Args:
+        name: the name of the TargetPlatformModel (not used in this function).
+        tp_model: a TargetPlatformModel to return.
+
+    Returns:
+        The given TargetPlatformModel object.
+
+    """
+
+    return tp_model

model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/latest/__init__.py

@@ -16,9 +16,10 @@ from model_compression_toolkit.verify_packages import FOUND_TORCH, FOUND_TF
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tp_model import get_tp_model, generate_tp_model, \
     get_op_quantization_configs
 if FOUND_TF:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_keras import get_keras_tpc as get_keras_tpc_latest
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_keras import generate_keras_tpc
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tp_model import get_tp_model as get_keras_tpc_latest
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.get_target_platform_capabilities import \
+        get_tpc_model as generate_keras_tpc
 if FOUND_TORCH:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import get_pytorch_tpc as \
-        get_pytorch_tpc_latest
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import generate_pytorch_tpc
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tp_model import get_tp_model as get_pytorch_tpc_latest
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.get_target_platform_capabilities import \
+        get_tpc_model as generate_pytorch_tpc

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

@@ -167,38 +167,70 @@ def generate_tp_model(default_config: OpQuantizationConfig,
     operator_set = []
     fusing_patterns = []
 
-    operator_set.append(schema.OperatorsSet(name="NoQuantization",
-                           qc_options=default_configuration_options.clone_and_edit(enable_activation_quantization=False)
-                           .clone_and_edit_weight_attribute(enable_weights_quantization=False)))
+    no_quantization_config = (default_configuration_options.clone_and_edit(enable_activation_quantization=False)
+                              .clone_and_edit_weight_attribute(enable_weights_quantization=False))
+
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.STACK, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.UNSTACK, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.DROPOUT, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.FLATTEN, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.SPLIT_CHUNK, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.GET_ITEM, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.RESHAPE, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.UNSQUEEZE, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.SIZE, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.PERMUTE, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.TRANSPOSE, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.EQUAL, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.ARGMAX, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.GATHER, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.TOPK, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.SQUEEZE, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.MAXPOOL, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.PAD, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.ZERO_PADDING2D, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.CAST, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.COMBINED_NON_MAX_SUPPRESSION, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.FAKE_QUANT, qc_options=no_quantization_config))
+    operator_set.append(schema.OperatorsSet(name=schema.OperatorSetNames.SSD_POST_PROCESS, qc_options=no_quantization_config))
 
     # Define operator sets that use mixed_precision_configuration_options:
-    conv = schema.OperatorsSet(name="Conv", qc_options=mixed_precision_configuration_options)
-    fc = schema.OperatorsSet(name="FullyConnected", qc_options=mixed_precision_configuration_options)
+    conv = schema.OperatorsSet(name=schema.OperatorSetNames.CONV, qc_options=mixed_precision_configuration_options)
+    conv_transpose = schema.OperatorsSet(name=schema.OperatorSetNames.CONV_TRANSPOSE, qc_options=mixed_precision_configuration_options)
+    depthwise_conv = schema.OperatorsSet(name=schema.OperatorSetNames.DEPTHWISE_CONV, qc_options=mixed_precision_configuration_options)
+    fc = schema.OperatorsSet(name=schema.OperatorSetNames.FULLY_CONNECTED, qc_options=mixed_precision_configuration_options)
 
     # Define operations sets without quantization configuration
     # options (useful for creating fusing patterns, for example):
-    any_relu = schema.OperatorsSet(name="AnyReLU")
-    add = schema.OperatorsSet(name="Add")
-    sub = schema.OperatorsSet(name="Sub")
-    mul = schema.OperatorsSet(name="Mul")
-    div = schema.OperatorsSet(name="Div")
-    prelu = schema.OperatorsSet(name="PReLU")
-    swish = schema.OperatorsSet(name="Swish")
-    sigmoid = schema.OperatorsSet(name="Sigmoid")
-    tanh = schema.OperatorsSet(name="Tanh")
-
-    operator_set.extend([conv, fc, any_relu, add, sub, mul, div, prelu, swish, sigmoid, tanh])
+    relu = schema.OperatorsSet(name=schema.OperatorSetNames.RELU)
+    relu6 = schema.OperatorsSet(name=schema.OperatorSetNames.RELU6)
+    leaky_relu = schema.OperatorsSet(name=schema.OperatorSetNames.LEAKY_RELU)
+    prelu = schema.OperatorsSet(name=schema.OperatorSetNames.PRELU)
+    add = schema.OperatorsSet(name=schema.OperatorSetNames.ADD)
+    sub = schema.OperatorsSet(name=schema.OperatorSetNames.SUB)
+    mul = schema.OperatorsSet(name=schema.OperatorSetNames.MUL)
+    div = schema.OperatorsSet(name=schema.OperatorSetNames.DIV)
+    swish = schema.OperatorsSet(name=schema.OperatorSetNames.SWISH)
+    hard_swish = schema.OperatorsSet(name=schema.OperatorSetNames.HARDSWISH)
+    sigmoid = schema.OperatorsSet(name=schema.OperatorSetNames.SIGMOID)
+    tanh = schema.OperatorsSet(name=schema.OperatorSetNames.TANH)
+    hard_tanh = schema.OperatorsSet(name=schema.OperatorSetNames.HARD_TANH)
+
+    operator_set.extend([conv, conv_transpose, depthwise_conv, fc, relu, relu6, leaky_relu, add, sub, mul, div, prelu, swish,
+                         hard_swish, sigmoid, tanh, hard_tanh])
+    any_relu = schema.OperatorSetConcat(operators_set=[relu, relu6, leaky_relu, hard_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.
     # To group multiple sets with regard to fusing, an OperatorSetConcat can be created
-    activations_after_conv_to_fuse = schema.OperatorSetConcat(operators_set=[any_relu, swish, prelu, sigmoid, tanh])
-    activations_after_fc_to_fuse = schema.OperatorSetConcat(operators_set=[any_relu, swish, sigmoid])
+    activations_after_conv_to_fuse = schema.OperatorSetConcat(operators_set=[relu, relu6, leaky_relu, hard_tanh, swish, hard_swish, prelu, sigmoid, tanh])
+    conv_types = schema.OperatorSetConcat(operators_set=[conv, conv_transpose, depthwise_conv])
+    activations_after_fc_to_fuse = schema.OperatorSetConcat(operators_set=[relu, relu6, leaky_relu, hard_tanh, swish, hard_swish, sigmoid])
     any_binary = schema.OperatorSetConcat(operators_set=[add, sub, mul, div])
 
     # ------------------- #
     # Fusions
     # ------------------- #
-    fusing_patterns.append(schema.Fusing(operator_groups=(conv, activations_after_conv_to_fuse)))
+    fusing_patterns.append(schema.Fusing(operator_groups=(conv_types, activations_after_conv_to_fuse)))
     fusing_patterns.append(schema.Fusing(operator_groups=(fc, activations_after_fc_to_fuse)))
     fusing_patterns.append(schema.Fusing(operator_groups=(any_binary, any_relu)))
 

model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/latest/__init__.py

@@ -15,8 +15,12 @@
 from model_compression_toolkit.verify_packages import FOUND_TORCH, FOUND_TF
 from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tp_model import get_tp_model, generate_tp_model, get_op_quantization_configs
 if FOUND_TF:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tpc_keras import get_keras_tpc as get_keras_tpc_latest
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tpc_keras import generate_keras_tpc
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tp_model import get_tp_model as \
+        get_keras_tpc_latest
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.get_target_platform_capabilities import \
+        get_tpc_model as generate_keras_tpc, get_tpc_model as generate_keras_tpc
 if FOUND_TORCH:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tpc_pytorch import get_pytorch_tpc as get_pytorch_tpc_latest
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tpc_pytorch import generate_pytorch_tpc
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.qnnpack_tpc.v1.tp_model import get_tp_model as \
+        get_pytorch_tpc_latest
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.get_target_platform_capabilities import \
+        get_tpc_model as generate_pytorch_tpc, get_tpc_model as generate_pytorch_tpc

model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/v1/tp_model.py

@@ -148,19 +148,29 @@ def generate_tp_model(default_config: OpQuantizationConfig,
     operator_set = []
     fusing_patterns = []
 
-    conv = schema.OperatorsSet(name="Conv")
-    batchnorm = schema.OperatorsSet(name="BatchNorm")
-    relu = schema.OperatorsSet(name="Relu")
-    linear = schema.OperatorsSet(name="Linear")
+    conv = schema.OperatorsSet(name=schema.OperatorSetNames.CONV)
+    conv_depthwise = schema.OperatorsSet(name=schema.OperatorSetNames.DEPTHWISE_CONV)
+    conv_transpose = schema.OperatorsSet(name=schema.OperatorSetNames.CONV_TRANSPOSE)
+    batchnorm = schema.OperatorsSet(name=schema.OperatorSetNames.BATCH_NORM)
+    relu = schema.OperatorsSet(name=schema.OperatorSetNames.RELU)
+    relu6 = schema.OperatorsSet(name=schema.OperatorSetNames.RELU6)
+
+    hard_tanh = schema.OperatorsSet(name=schema.OperatorSetNames.HARD_TANH)
+    linear = schema.OperatorsSet(name=schema.OperatorSetNames.FULLY_CONNECTED)
+
+    operator_set.extend([conv, conv_depthwise, conv_transpose, batchnorm, relu, relu6, hard_tanh, linear])
+
+    conv_opset_concat = schema.OperatorSetConcat(operators_set=[conv, conv_transpose])
+    relu_opset_concat = schema.OperatorSetConcat(operators_set=[relu, relu6, hard_tanh])
 
-    operator_set.extend([conv, batchnorm, relu, linear])
     # ------------------- #
     # Fusions
     # ------------------- #
-    fusing_patterns.append(schema.Fusing(operator_groups=(conv, batchnorm, relu)))
-    fusing_patterns.append(schema.Fusing(operator_groups=(conv, batchnorm)))
-    fusing_patterns.append(schema.Fusing(operator_groups=(conv, relu)))
-    fusing_patterns.append(schema.Fusing(operator_groups=(linear, relu)))
+    fusing_patterns.append(schema.Fusing(operator_groups=(conv_opset_concat, batchnorm, relu_opset_concat)))
+    fusing_patterns.append(schema.Fusing(operator_groups=(conv_opset_concat, batchnorm)))
+    fusing_patterns.append(schema.Fusing(operator_groups=(conv_opset_concat, relu_opset_concat)))
+    fusing_patterns.append(schema.Fusing(operator_groups=(linear, relu_opset_concat)))
+
     # 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):

model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/latest/__init__.py

@@ -15,8 +15,11 @@
 from model_compression_toolkit.verify_packages import FOUND_TORCH, FOUND_TF
 from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tp_model import get_tp_model, generate_tp_model, get_op_quantization_configs
 if FOUND_TF:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tpc_keras import get_keras_tpc as get_keras_tpc_latest
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tpc_keras import generate_keras_tpc
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tp_model import get_keras_tpc as get_keras_tpc_latest
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.get_target_platform_capabilities import \
+        get_tpc_model as generate_keras_tpc
 if FOUND_TORCH:
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tpc_pytorch import get_pytorch_tpc as get_pytorch_tpc_latest
-    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tpc_pytorch import generate_pytorch_tpc
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.tflite_tpc.v1.tp_model import \
+        get_tp_model as get_pytorch_tpc_latest
+    from model_compression_toolkit.target_platform_capabilities.tpc_models.get_target_platform_capabilities import \
+        get_tpc_model as generate_pytorch_tpc