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

model_compression_toolkit/pruning/keras/pruning_facade.py

@@ -17,12 +17,12 @@ from typing import Callable, Tuple
 
 from model_compression_toolkit import get_target_platform_capabilities
 from model_compression_toolkit.constants import TENSORFLOW
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.pruning.pruner import Pruner
 from model_compression_toolkit.core.common.pruning.pruning_config import PruningConfig
 from model_compression_toolkit.core.common.pruning.pruning_info import PruningInfo
-from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.set_node_quantization_config import set_quantization_configuration_to_graph
 from model_compression_toolkit.core.graph_prep_runner import read_model_to_graph
 from model_compression_toolkit.logger import Logger
@@ -35,6 +35,8 @@ if FOUND_TF:
     from model_compression_toolkit.core.keras.pruning.pruning_keras_implementation import PruningKerasImplementation
     from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from tensorflow.keras.models import Model
+    from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework.attach2keras import \
+        AttachTpcToKeras
 
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
@@ -42,7 +44,7 @@ if FOUND_TF:
                                    target_resource_utilization: ResourceUtilization,
                                    representative_data_gen: Callable,
                                    pruning_config: PruningConfig = PruningConfig(),
-                                   target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> Tuple[Model, PruningInfo]:
+                                   target_platform_capabilities: TargetPlatformModel = DEFAULT_KERAS_TPC) -> Tuple[Model, PruningInfo]:
         """
         Perform structured pruning on a Keras model to meet a specified target resource utilization.
         This function prunes the provided model according to the target resource utilization by grouping and pruning
@@ -111,6 +113,10 @@ if FOUND_TF:
         # Instantiate the Keras framework implementation.
         fw_impl = PruningKerasImplementation()
 
+        # Attach tpc model to framework
+        attach2keras = AttachTpcToKeras()
+        target_platform_capabilities = attach2keras.attach(target_platform_capabilities)
+
         # Convert the original Keras model to an internal graph representation.
         float_graph = read_model_to_graph(model,
                                           representative_data_gen,

model_compression_toolkit/pruning/pytorch/pruning_facade.py

@@ -16,12 +16,12 @@
 from typing import Callable, Tuple
 from model_compression_toolkit import get_target_platform_capabilities
 from model_compression_toolkit.constants import PYTORCH
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.pruning.pruner import Pruner
 from model_compression_toolkit.core.common.pruning.pruning_config import PruningConfig
 from model_compression_toolkit.core.common.pruning.pruning_info import PruningInfo
-from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.set_node_quantization_config import set_quantization_configuration_to_graph
 from model_compression_toolkit.core.graph_prep_runner import read_model_to_graph
 from model_compression_toolkit.logger import Logger
@@ -38,6 +38,8 @@ if FOUND_TORCH:
         PruningPytorchImplementation
     from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
     from torch.nn import Module
+    from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework.attach2pytorch import \
+        AttachTpcToPytorch
 
     # Set the default Target Platform Capabilities (TPC) for PyTorch.
     DEFAULT_PYOTRCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
@@ -46,7 +48,7 @@ if FOUND_TORCH:
                                      target_resource_utilization: ResourceUtilization,
                                      representative_data_gen: Callable,
                                      pruning_config: PruningConfig = PruningConfig(),
-                                     target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYOTRCH_TPC) -> \
+                                     target_platform_capabilities: TargetPlatformModel = DEFAULT_PYOTRCH_TPC) -> \
             Tuple[Module, PruningInfo]:
         """
         Perform structured pruning on a Pytorch model to meet a specified target resource utilization.
@@ -117,6 +119,10 @@ if FOUND_TORCH:
         # Instantiate the Pytorch framework implementation.
         fw_impl = PruningPytorchImplementation()
 
+        # Attach TPC to framework
+        attach2pytorch = AttachTpcToPytorch()
+        target_platform_capabilities = attach2pytorch.attach(target_platform_capabilities)
+
         # Convert the original Pytorch model to an internal graph representation.
         float_graph = read_model_to_graph(model,
                                           representative_data_gen,

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -22,6 +22,7 @@ from model_compression_toolkit.core.common.quantization.quantize_graph_weights i
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
@@ -41,6 +42,9 @@ if FOUND_TF:
 
     from model_compression_toolkit import get_target_platform_capabilities
     from mct_quantizers.keras.metadata import add_metadata
+    from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework.attach2keras import \
+        AttachTpcToKeras
+
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
 
@@ -48,7 +52,7 @@ if FOUND_TF:
                                          representative_data_gen: Callable,
                                          target_resource_utilization: ResourceUtilization = None,
                                          core_config: CoreConfig = CoreConfig(),
-                                         target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC):
+                                         target_platform_capabilities: TargetPlatformModel = DEFAULT_KERAS_TPC):
         """
          Quantize a trained Keras model using post-training quantization. The model is quantized using a
          symmetric constraint quantization thresholds (power of two).
@@ -134,6 +138,11 @@ if FOUND_TF:
 
         fw_impl = KerasImplementation()
 
+        attach2keras = AttachTpcToKeras()
+        target_platform_capabilities = attach2keras.attach(
+            target_platform_capabilities,
+            custom_opset2layer=core_config.quantization_config.custom_tpc_opset_to_layer)
+
         # Ignore returned hessian service as PTQ does not use it
         tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_model,
                                                                 representative_data_gen=representative_data_gen,

model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -19,6 +19,7 @@ from typing import Callable
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import PYTORCH
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
@@ -39,6 +40,8 @@ if FOUND_TORCH:
     from model_compression_toolkit.exporter.model_wrapper.pytorch.builder.fully_quantized_model_builder import get_exportable_pytorch_model
     from model_compression_toolkit import get_target_platform_capabilities
     from mct_quantizers.pytorch.metadata import add_metadata
+    from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework.attach2pytorch import \
+        AttachTpcToPytorch
 
     DEFAULT_PYTORCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
 
@@ -46,7 +49,7 @@ if FOUND_TORCH:
                                            representative_data_gen: Callable,
                                            target_resource_utilization: ResourceUtilization = None,
                                            core_config: CoreConfig = CoreConfig(),
-                                           target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
+                                           target_platform_capabilities: TargetPlatformModel = DEFAULT_PYTORCH_TPC):
         """
         Quantize a trained Pytorch module using post-training quantization.
         By default, the module is quantized using a symmetric constraint quantization thresholds
@@ -107,6 +110,11 @@ if FOUND_TORCH:
 
         fw_impl = PytorchImplementation()
 
+        # Attach tpc model to framework
+        attach2pytorch = AttachTpcToPytorch()
+        target_platform_capabilities = attach2pytorch.attach(target_platform_capabilities,
+                                                             core_config.quantization_config.custom_tpc_opset_to_layer)
+
         # Ignore hessian info service as it is not used here yet.
         tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_module,
                                                                 representative_data_gen=representative_data_gen,

model_compression_toolkit/qat/__init__.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 from model_compression_toolkit.qat.common.qat_config import QATConfig
+from model_compression_toolkit.verify_packages import FOUND_TF, FOUND_TORCH
 
-from model_compression_toolkit.qat.keras.quantization_facade import keras_quantization_aware_training_init_experimental, keras_quantization_aware_training_finalize_experimental
-from model_compression_toolkit.qat.pytorch.quantization_facade import pytorch_quantization_aware_training_init_experimental, pytorch_quantization_aware_training_finalize_experimental
+if FOUND_TF:
+    from model_compression_toolkit.qat.keras.quantization_facade import keras_quantization_aware_training_init_experimental, keras_quantization_aware_training_finalize_experimental
+if FOUND_TORCH:
+    from model_compression_toolkit.qat.pytorch.quantization_facade import pytorch_quantization_aware_training_init_experimental, pytorch_quantization_aware_training_finalize_experimental

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -19,6 +19,7 @@ from functools import partial
 from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
@@ -54,6 +55,8 @@ if FOUND_TF:
     from model_compression_toolkit.qat.keras.quantizer.quantization_builder import quantization_builder, \
     get_activation_quantizer_holder
     from model_compression_toolkit.qat.common.qat_config import QATConfig
+    from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework.attach2keras import \
+        AttachTpcToKeras
 
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
@@ -90,7 +93,7 @@ if FOUND_TF:
                                                             target_resource_utilization: ResourceUtilization = None,
                                                             core_config: CoreConfig = CoreConfig(),
                                                             qat_config: QATConfig = QATConfig(),
-                                                            target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC):
+                                                            target_platform_capabilities: TargetPlatformModel = DEFAULT_KERAS_TPC):
         """
          Prepare a trained Keras model for quantization aware training. First the model quantization is optimized
          with post-training quantization, then the model layers are wrapped with QuantizeWrappers. The model is
@@ -186,6 +189,11 @@ if FOUND_TF:
 
         fw_impl = KerasImplementation()
 
+        attach2keras = AttachTpcToKeras()
+        target_platform_capabilities = attach2keras.attach(
+            target_platform_capabilities,
+            custom_opset2layer=core_config.quantization_config.custom_tpc_opset_to_layer)
+
         # Ignore hessian service since is not used in QAT at the moment
         tg, bit_widths_config, _, _ = core_runner(in_model=in_model,
                                                   representative_data_gen=representative_data_gen,

model_compression_toolkit/qat/pytorch/quantization_facade.py

@@ -17,6 +17,9 @@ from typing import Callable
 from functools import partial
 
 from model_compression_toolkit.constants import PYTORCH
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
+from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework.attach2pytorch import \
+    AttachTpcToPytorch
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 
 from model_compression_toolkit.core import CoreConfig
@@ -79,7 +82,7 @@ if FOUND_TORCH:
                                                               target_resource_utilization: ResourceUtilization = None,
                                                               core_config: CoreConfig = CoreConfig(),
                                                               qat_config: QATConfig = QATConfig(),
-                                                              target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
+                                                              target_platform_capabilities: TargetPlatformModel = DEFAULT_PYTORCH_TPC):
         """
          Prepare a trained Pytorch model for quantization aware training. First the model quantization is optimized
          with post-training quantization, then the model layers are wrapped with QuantizeWrappers. The model is
@@ -154,6 +157,11 @@ if FOUND_TORCH:
         tb_w = init_tensorboard_writer(DEFAULT_PYTORCH_INFO)
         fw_impl = PytorchImplementation()
 
+        # Attach tpc model to framework
+        attach2pytorch = AttachTpcToPytorch()
+        target_platform_capabilities = attach2pytorch.attach(target_platform_capabilities,
+                                                             core_config.quantization_config.custom_tpc_opset_to_layer)
+
         # Ignore hessian scores service as we do not use it here
         tg, bit_widths_config, _, _ = core_runner(in_model=in_model,
                                                   representative_data_gen=representative_data_gen,

model_compression_toolkit/target_platform_capabilities/schema/mct_current_schema.py

@@ -7,6 +7,6 @@ OpQuantizationConfig = schema.OpQuantizationConfig
 QuantizationConfigOptions = schema.QuantizationConfigOptions
 OperatorsSetBase = schema.OperatorsSetBase
 OperatorsSet = schema.OperatorsSet
-OperatorSetConcat= schema.OperatorSetConcat
+OperatorSetConcat = schema.OperatorSetConcat
 Fusing = schema.Fusing
 TargetPlatformModel = schema.TargetPlatformModel

model_compression_toolkit/target_platform_capabilities/schema/v1.py

@@ -13,66 +13,74 @@
 # limitations under the License.
 # ==============================================================================
 import pprint
-
 from enum import Enum
 from typing import Dict, Any, Union, Tuple, List, Optional, Literal, Annotated
+
+from pydantic import BaseModel, Field, root_validator, validator, PositiveInt
+
 from mct_quantizers import QuantizationMethod
 from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.logger import Logger
-from pydantic import BaseModel, Field, root_validator, validator, PositiveInt, PrivateAttr
-
-
-class OperatorSetNames(Enum):
-    OPSET_CONV = "Conv"
-    OPSET_DEPTHWISE_CONV = "DepthwiseConv2D"
-    OPSET_CONV_TRANSPOSE = "ConvTranspose"
-    OPSET_FULLY_CONNECTED = "FullyConnected"
-    OPSET_CONCATENATE = "Concatenate"
-    OPSET_STACK = "Stack"
-    OPSET_UNSTACK = "Unstack"
-    OPSET_GATHER = "Gather"
-    OPSET_EXPAND = "Expend"
-    OPSET_BATCH_NORM = "BatchNorm"
-    OPSET_RELU = "ReLU"
-    OPSET_RELU6 = "ReLU6"
-    OPSET_LEAKY_RELU = "LEAKYReLU"
-    OPSET_HARD_TANH = "HardTanh"
-    OPSET_ADD = "Add"
-    OPSET_SUB = "Sub"
-    OPSET_MUL = "Mul"
-    OPSET_DIV = "Div"
-    OPSET_MIN = "Min"
-    OPSET_MAX = "Max"
-    OPSET_PRELU = "PReLU"
-    OPSET_SWISH = "Swish"
-    OPSET_SIGMOID = "Sigmoid"
-    OPSET_TANH = "Tanh"
-    OPSET_GELU = "Gelu"
-    OPSET_HARDSIGMOID = "HardSigmoid"
-    OPSET_HARDSWISH = "HardSwish"
-    OPSET_FLATTEN = "Flatten"
-    OPSET_GET_ITEM = "GetItem"
-    OPSET_RESHAPE = "Reshape"
-    OPSET_UNSQUEEZE = "Unsqueeze"
-    OPSET_SQUEEZE = "Squeeze"
-    OPSET_PERMUTE = "Permute"
-    OPSET_TRANSPOSE = "Transpose"
-    OPSET_DROPOUT = "Dropout"
-    OPSET_SPLIT = "Split"
-    OPSET_CHUNK = "Chunk"
-    OPSET_MAXPOOL = "MaxPool"
-    OPSET_SIZE = "Size"
-    OPSET_SHAPE = "Shape"
-    OPSET_EQUAL = "Equal"
-    OPSET_ARGMAX = "ArgMax"
-    OPSET_TOPK = "TopK"
-    OPSET_FAKE_QUANT_WITH_MIN_MAX_VARS = "FakeQuantWithMinMaxVars"
-    OPSET_COMBINED_NON_MAX_SUPPRESSION = "CombinedNonMaxSuppression"
-    OPSET_CROPPING2D = "Cropping2D"
-    OPSET_ZERO_PADDING2d = "ZeroPadding2D"
-    OPSET_CAST = "Cast"
-    OPSET_STRIDED_SLICE = "StridedSlice"
-    OPSET_SSD_POST_PROCESS = "SSDPostProcess"
+
+
+class OperatorSetNames(str, Enum):
+    CONV = "Conv"
+    DEPTHWISE_CONV = "DepthwiseConv2D"
+    CONV_TRANSPOSE = "ConvTranspose"
+    FULLY_CONNECTED = "FullyConnected"
+    CONCATENATE = "Concatenate"
+    STACK = "Stack"
+    UNSTACK = "Unstack"
+    GATHER = "Gather"
+    EXPAND = "Expend"
+    BATCH_NORM = "BatchNorm"
+    L2NORM = "L2Norm"
+    RELU = "ReLU"
+    RELU6 = "ReLU6"
+    LEAKY_RELU = "LeakyReLU"
+    ELU = "Elu"
+    HARD_TANH = "HardTanh"
+    ADD = "Add"
+    SUB = "Sub"
+    MUL = "Mul"
+    DIV = "Div"
+    MIN = "Min"
+    MAX = "Max"
+    PRELU = "PReLU"
+    ADD_BIAS = "AddBias"
+    SWISH = "Swish"
+    SIGMOID = "Sigmoid"
+    SOFTMAX = "Softmax"
+    LOG_SOFTMAX = "LogSoftmax"
+    TANH = "Tanh"
+    GELU = "Gelu"
+    HARDSIGMOID = "HardSigmoid"
+    HARDSWISH = "HardSwish"
+    FLATTEN = "Flatten"
+    GET_ITEM = "GetItem"
+    RESHAPE = "Reshape"
+    UNSQUEEZE = "Unsqueeze"
+    SQUEEZE = "Squeeze"
+    PERMUTE = "Permute"
+    TRANSPOSE = "Transpose"
+    DROPOUT = "Dropout"
+    SPLIT_CHUNK = "SplitChunk"
+    MAXPOOL = "MaxPool"
+    AVGPOOL = "AvgPool"
+    SIZE = "Size"
+    SHAPE = "Shape"
+    EQUAL = "Equal"
+    ARGMAX = "ArgMax"
+    TOPK = "TopK"
+    FAKE_QUANT = "FakeQuant"
+    COMBINED_NON_MAX_SUPPRESSION = "CombinedNonMaxSuppression"
+    ZERO_PADDING2D = "ZeroPadding2D"
+    CAST = "Cast"
+    RESIZE = "Resize"
+    PAD = "Pad"
+    FOLD = "Fold"
+    STRIDED_SLICE = "StridedSlice"
+    SSD_POST_PROCESS = "SSDPostProcess"
 
     @classmethod
     def get_values(cls):

model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/attach2fw.py

@@ -1,12 +1,15 @@
-from typing import Dict, Tuple, List, Any, Optional
+from typing import Dict, Optional
 
-from model_compression_toolkit import DefaultDict
-from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel, \
+    OperatorsSet
 from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities, \
     OperationsSetToLayers
 
+from model_compression_toolkit.core.common.quantization.quantization_config import CustomOpsetLayers
 
-class AttachTpModelToFw:
+
+class AttachTpcToFramework:
 
     def __init__(self):
         self._opset2layer = None
@@ -17,7 +20,7 @@ class AttachTpModelToFw:
         self._opset2attr_mapping = None  # Mapping of operation sets to their corresponding framework-specific layers
 
     def attach(self, tpc_model: TargetPlatformModel,
-               custom_opset2layer: Dict[str, Tuple[List[Any], Optional[Dict[str, DefaultDict]]]] = None
+               custom_opset2layer: Optional[Dict[str, 'CustomOpsetLayers']] = None
                ) -> TargetPlatformCapabilities:
         """
         Attaching a TargetPlatformModel which includes a platform capabilities description to specific
@@ -35,22 +38,30 @@ class AttachTpModelToFw:
         """
 
         tpc = TargetPlatformCapabilities(tpc_model)
+        custom_opset2layer = custom_opset2layer if custom_opset2layer is not None else {}
 
         with tpc:
-            for opset_name, operators in self._opset2layer.items():
-                attr_mapping = self._opset2attr_mapping.get(opset_name)
-                OperationsSetToLayers(opset_name, operators, attr_mapping=attr_mapping)
-
-            if custom_opset2layer is not None:
-                for opset_name, operators in custom_opset2layer.items():
-                    if len(operators) == 1:
-                        OperationsSetToLayers(opset_name, operators[0])
-                    elif len(operators) == 2:
-                        OperationsSetToLayers(opset_name, operators[0], attr_mapping=operators[1])
+            for opset in tpc_model.operator_set:
+                if isinstance(opset, OperatorsSet):  # filter out OperatorsSetConcat
+                    if opset.name in custom_opset2layer:
+                        custom_opset_layers = custom_opset2layer[opset.name]
+                        OperationsSetToLayers(opset.name,
+                                              layers=custom_opset_layers.operators,
+                                              attr_mapping=custom_opset_layers.attr_mapping)
+
+                    elif opset.name in self._opset2layer:
+                        # Note that if the user provided a custom operator set with a name that exists in our
+                        # pre-defined set of operator sets, we prioritize the user's custom opset definition
+                        layers = self._opset2layer[opset.name]
+                        if len(layers) > 0:
+                            # If the framework does not define any matching operators to a given operator set name that
+                            # appears in the TPC, then we just skip it
+                            attr_mapping = self._opset2attr_mapping.get(opset.name)
+                            OperationsSetToLayers(opset.name, layers, attr_mapping=attr_mapping)
                     else:
-                        raise ValueError(f"Custom operator set to layer mapping should include up to 2 elements - "
-                                         f"a list of layers to attach to the operator and an optional mapping of "
-                                         f"attributes names, but given a mapping contains {len(operators)} elements.")
+                        Logger.critical(f'{opset.name} is defined in TargetPlatformModel, '
+                                        f'but is not defined in the framework set of operators or in the provided '
+                                        f'custom operator sets mapping.')
 
         return tpc
 

model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/attach2keras.py

@@ -23,12 +23,12 @@ if FOUND_SONY_CUSTOM_LAYERS:
 
 if version.parse(tf.__version__) >= version.parse("2.13"):
     from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
-        MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
-        Conv2DTranspose, Identity, Concatenate, BatchNormalization, Minimum, Maximum
+        MaxPooling2D, AveragePooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
+        Conv2DTranspose, Concatenate, BatchNormalization, Minimum, Maximum, Softmax
 else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
-        MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
-        Conv2DTranspose, Concatenate, BatchNormalization, Minimum, Maximum
+        MaxPooling2D, AveragePooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
+        Conv2DTranspose, Concatenate, BatchNormalization, Minimum, Maximum, Softmax
 
 from model_compression_toolkit import DefaultDict
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, BIAS, \
@@ -36,72 +36,93 @@ from model_compression_toolkit.target_platform_capabilities.constants import KER
 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.targetplatform2framework.attach2fw import \
-    AttachTpModelToFw
+    AttachTpcToFramework
 
 
-class AttachTpModelToKeras(AttachTpModelToFw):
+class AttachTpcToKeras(AttachTpcToFramework):
     def __init__(self):
         super().__init__()
 
         self._opset2layer = {
-            OperatorSetNames.OPSET_CONV.value: [Conv2D, tf.nn.conv2d],
-            OperatorSetNames.OPSET_DEPTHWISE_CONV.value: [DepthwiseConv2D, tf.nn.depthwise_conv2d],
-            OperatorSetNames.OPSET_CONV_TRANSPOSE.value: [Conv2DTranspose, tf.nn.conv2d_transpose],
-            OperatorSetNames.OPSET_FULLY_CONNECTED.value: [Dense],
-            OperatorSetNames.OPSET_CONCATENATE.value: [tf.concat, Concatenate],
-            OperatorSetNames.OPSET_STACK.value: [tf.stack],
-            OperatorSetNames.OPSET_UNSTACK.value: [tf.unstack],
-            OperatorSetNames.OPSET_GATHER.value: [tf.gather, tf.compat.v1.gather],
-            OperatorSetNames.OPSET_EXPAND.value: [],
-            OperatorSetNames.OPSET_BATCH_NORM.value: [BatchNormalization],
-            OperatorSetNames.OPSET_RELU.value: [tf.nn.relu, ReLU],
-            OperatorSetNames.OPSET_RELU6.value: [tf.nn.relu6],
-            OperatorSetNames.OPSET_LEAKY_RELU.value: [tf.nn.leaky_relu, LeakyReLU],
-            OperatorSetNames.OPSET_HARD_TANH.value: [LayerFilterParams(Activation, activation="hard_tanh")],
-            OperatorSetNames.OPSET_ADD.value: [tf.add, Add],
-            OperatorSetNames.OPSET_SUB.value: [tf.subtract, Subtract],
-            OperatorSetNames.OPSET_MUL.value: [tf.math.multiply, Multiply],
-            OperatorSetNames.OPSET_DIV.value: [tf.math.divide, tf.math.truediv],
-            OperatorSetNames.OPSET_MIN.value: [tf.math.minimum, Minimum],
-            OperatorSetNames.OPSET_MAX.value: [tf.math.maximum, Maximum],
-            OperatorSetNames.OPSET_PRELU.value: [PReLU],
-            OperatorSetNames.OPSET_SWISH.value: [tf.nn.swish, LayerFilterParams(Activation, activation="swish")],
-            OperatorSetNames.OPSET_SIGMOID.value: [tf.nn.sigmoid, LayerFilterParams(Activation, activation="sigmoid")],
-            OperatorSetNames.OPSET_TANH.value: [tf.nn.tanh, LayerFilterParams(Activation, activation="tanh")],
-            OperatorSetNames.OPSET_GELU.value: [tf.nn.gelu, LayerFilterParams(Activation, activation="gelu")],
-            OperatorSetNames.OPSET_HARDSIGMOID.value: [tf.keras.activations.hard_sigmoid,
-                                                       LayerFilterParams(Activation, activation="hard_sigmoid")],
-            OperatorSetNames.OPSET_FLATTEN.value: [Flatten],
-            OperatorSetNames.OPSET_GET_ITEM.value: [tf.__operators__.getitem],
-            OperatorSetNames.OPSET_RESHAPE.value: [Reshape, tf.reshape],
-            OperatorSetNames.OPSET_PERMUTE.value: [Permute],
-            OperatorSetNames.OPSET_TRANSPOSE.value: [tf.transpose],
-            OperatorSetNames.OPSET_DROPOUT.value: [Dropout],
-            OperatorSetNames.OPSET_SPLIT.value: [tf.split],
-            OperatorSetNames.OPSET_MAXPOOL.value: [MaxPooling2D],
-            OperatorSetNames.OPSET_SHAPE.value: [tf.shape, tf.compat.v1.shape],
-            OperatorSetNames.OPSET_EQUAL.value: [tf.math.equal],
-            OperatorSetNames.OPSET_ARGMAX.value: [tf.math.argmax],
-            OperatorSetNames.OPSET_TOPK.value: [tf.nn.top_k],
-            OperatorSetNames.OPSET_FAKE_QUANT_WITH_MIN_MAX_VARS.value: [tf.quantization.fake_quant_with_min_max_vars],
-            OperatorSetNames.OPSET_COMBINED_NON_MAX_SUPPRESSION.value: [tf.image.combined_non_max_suppression],
-            OperatorSetNames.OPSET_CROPPING2D.value: [Cropping2D],
-            OperatorSetNames.OPSET_ZERO_PADDING2d.value: [ZeroPadding2D],
-            OperatorSetNames.OPSET_CAST.value: [tf.cast],
-            OperatorSetNames.OPSET_STRIDED_SLICE.value: [tf.strided_slice]
+            OperatorSetNames.CONV: [Conv2D, tf.nn.conv2d],
+            OperatorSetNames.DEPTHWISE_CONV: [DepthwiseConv2D, tf.nn.depthwise_conv2d],
+            OperatorSetNames.CONV_TRANSPOSE: [Conv2DTranspose, tf.nn.conv2d_transpose],
+            OperatorSetNames.FULLY_CONNECTED: [Dense],
+            OperatorSetNames.CONCATENATE: [tf.concat, Concatenate],
+            OperatorSetNames.STACK: [tf.stack],
+            OperatorSetNames.UNSTACK: [tf.unstack],
+            OperatorSetNames.GATHER: [tf.gather, tf.compat.v1.gather],
+            OperatorSetNames.EXPAND: [],
+            OperatorSetNames.BATCH_NORM: [BatchNormalization, tf.nn.batch_normalization],
+            OperatorSetNames.RELU: [tf.nn.relu, ReLU, LayerFilterParams(Activation, activation="relu")],
+            OperatorSetNames.RELU6: [tf.nn.relu6],
+            OperatorSetNames.LEAKY_RELU: [tf.nn.leaky_relu, LeakyReLU, LayerFilterParams(Activation, activation="leaky_relu")],
+            OperatorSetNames.HARD_TANH: [LayerFilterParams(Activation, activation="hard_tanh")],
+            OperatorSetNames.ADD: [tf.add, Add],
+            OperatorSetNames.SUB: [tf.subtract, Subtract],
+            OperatorSetNames.MUL: [tf.math.multiply, Multiply],
+            OperatorSetNames.DIV: [tf.math.divide, tf.math.truediv],
+            OperatorSetNames.MIN: [tf.math.minimum, Minimum],
+            OperatorSetNames.MAX: [tf.math.maximum, Maximum],
+            OperatorSetNames.PRELU: [PReLU],
+            OperatorSetNames.SWISH: [tf.nn.swish, LayerFilterParams(Activation, activation="swish")],
+            OperatorSetNames.HARDSWISH: [LayerFilterParams(Activation, activation="hard_swish")],
+            OperatorSetNames.SIGMOID: [tf.nn.sigmoid, LayerFilterParams(Activation, activation="sigmoid")],
+            OperatorSetNames.TANH: [tf.nn.tanh, LayerFilterParams(Activation, activation="tanh")],
+            OperatorSetNames.GELU: [tf.nn.gelu, LayerFilterParams(Activation, activation="gelu")],
+            OperatorSetNames.HARDSIGMOID: [tf.keras.activations.hard_sigmoid,
+                                           LayerFilterParams(Activation, activation="hard_sigmoid")],
+            OperatorSetNames.FLATTEN: [Flatten],
+            OperatorSetNames.GET_ITEM: [tf.__operators__.getitem],
+            OperatorSetNames.RESHAPE: [Reshape, tf.reshape],
+            OperatorSetNames.PERMUTE: [Permute],
+            OperatorSetNames.TRANSPOSE: [tf.transpose],
+            OperatorSetNames.UNSQUEEZE: [tf.expand_dims],
+            OperatorSetNames.SQUEEZE: [tf.squeeze],
+            OperatorSetNames.DROPOUT: [Dropout],
+            OperatorSetNames.SPLIT_CHUNK: [tf.split],
+            OperatorSetNames.MAXPOOL: [MaxPooling2D, tf.nn.avg_pool2d],
+            OperatorSetNames.AVGPOOL: [AveragePooling2D],
+            OperatorSetNames.SIZE: [tf.size],
+            OperatorSetNames.RESIZE: [tf.image.resize],
+            OperatorSetNames.PAD: [tf.pad, Cropping2D],
+            OperatorSetNames.FOLD: [tf.space_to_batch_nd],
+            OperatorSetNames.SHAPE: [tf.shape, tf.compat.v1.shape],
+            OperatorSetNames.EQUAL: [tf.math.equal],
+            OperatorSetNames.ARGMAX: [tf.math.argmax],
+            OperatorSetNames.TOPK: [tf.nn.top_k],
+            OperatorSetNames.FAKE_QUANT: [tf.quantization.fake_quant_with_min_max_vars],
+            OperatorSetNames.COMBINED_NON_MAX_SUPPRESSION: [tf.image.combined_non_max_suppression],
+            OperatorSetNames.ZERO_PADDING2D: [ZeroPadding2D],
+            OperatorSetNames.CAST: [tf.cast],
+            OperatorSetNames.STRIDED_SLICE: [tf.strided_slice],
+            OperatorSetNames.ELU: [tf.nn.elu, LayerFilterParams(Activation, activation="elu")],
+            OperatorSetNames.SOFTMAX: [tf.nn.softmax, Softmax,
+                                       LayerFilterParams(Activation, activation="softmax")],
+            OperatorSetNames.LOG_SOFTMAX: [tf.nn.log_softmax],
+            OperatorSetNames.ADD_BIAS: [tf.nn.bias_add],
+            OperatorSetNames.L2NORM: [tf.math.l2_normalize],
         }
 
         if FOUND_SONY_CUSTOM_LAYERS:
-            self._opset2layer[OperatorSetNames.OPSET_POST_PROCESS] = [SSDPostProcess]
+            self._opset2layer[OperatorSetNames.SSD_POST_PROCESS] = [SSDPostProcess]
+        else:
+            # If Custom layers is not installed then we don't want the user to fail, but just ignore custom layers
+            # in the initialized framework TPC
+            self._opset2layer[OperatorSetNames.SSD_POST_PROCESS] = []
 
-        self._opset2attr_mapping = {OperatorSetNames.OPSET_CONV.value: {
-            KERNEL_ATTR: DefaultDict(default_value=KERAS_KERNEL),
-            BIAS_ATTR: DefaultDict(default_value=BIAS)},
-            OperatorSetNames.OPSET_DEPTHWISE_CONV.value: {
+        self._opset2attr_mapping = {
+            OperatorSetNames.CONV: {
+                KERNEL_ATTR: DefaultDict(default_value=KERAS_KERNEL),
+                BIAS_ATTR: DefaultDict(default_value=BIAS)},
+            OperatorSetNames.CONV_TRANSPOSE: {
+                KERNEL_ATTR: DefaultDict(default_value=KERAS_KERNEL),
+                BIAS_ATTR: DefaultDict(default_value=BIAS)},
+            OperatorSetNames.DEPTHWISE_CONV: {
                 KERNEL_ATTR: DefaultDict({
                     DepthwiseConv2D: KERAS_DEPTHWISE_KERNEL,
                     tf.nn.depthwise_conv2d: KERAS_DEPTHWISE_KERNEL}, default_value=KERAS_KERNEL),
                 BIAS_ATTR: DefaultDict(default_value=BIAS)},
-            OperatorSetNames.OPSET_FULLY_CONNECTED.value: {
+            OperatorSetNames.FULLY_CONNECTED: {
                 KERNEL_ATTR: DefaultDict(default_value=KERAS_KERNEL),
                 BIAS_ATTR: DefaultDict(default_value=BIAS)}}