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

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)}}

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