mct-nightly 2.2.0.20241119.516__py3-none-any.whl → 2.2.0.20241121.524__py3-none-any.whl

{mct_nightly-2.2.0.20241119.516.dist-info → mct_nightly-2.2.0.20241121.524.dist-info}/METADATA

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

{mct_nightly-2.2.0.20241119.516.dist-info → mct_nightly-2.2.0.20241121.524.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=LtqmejEQcQLXceDuGvAmjezjSZB8czuS4kfbMbUzF7k,1573
+model_compression_toolkit/__init__.py,sha256=X-Df1ZQGO1n9pF09XDXYplM7M5oQzCDKFyOASqQtRVA,1573
 model_compression_toolkit/constants.py,sha256=i4wYheBkIdQmsQA-axIpcT3YiSO1USNc-jaNiNE8w6E,3920
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -223,7 +223,7 @@ model_compression_toolkit/core/pytorch/constants.py,sha256=YwD_joIF0vK8UG2vW1NVv
 model_compression_toolkit/core/pytorch/data_util.py,sha256=YYbT135HhlTt0q6XdD2JX7AS_L92f_uV2rWq2hsJOCA,6325
 model_compression_toolkit/core/pytorch/default_framework_info.py,sha256=-Vls1P_8Ckm_18nnOsmQkZ71SmzHwtQLbQ383Z4Rb-U,4365
 model_compression_toolkit/core/pytorch/pytorch_device_config.py,sha256=S25cuw10AW3SEN_fRAGRcG_I3wdvvQx1ehSJzPnn-UI,4404
-model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=4uzO-lXfuitlC3NHx5-k2Fjm8VHa1T7ox9c8DSxYs9M,29437
+model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=SgxmSdzAQOPI9YHt4Q9-OeDi8fzAdgASHQ4nZ5maPsg,29599
 model_compression_toolkit/core/pytorch/pytorch_node_prior_info.py,sha256=2LDQ7qupglHQ7o1Am7LWdfYVacfQnl-aW2N6l9det1w,3264
 model_compression_toolkit/core/pytorch/resource_utilization_data_facade.py,sha256=xpKj99OZKT9NT0vKIl_cOe8d89d2gef1gKoNT6PFElE,4989
 model_compression_toolkit/core/pytorch/utils.py,sha256=7VbgcLwtQvdEEc_AJgSOQ3U3KRKCICFPaBirN1fIQxg,3940
@@ -246,6 +246,7 @@ model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/concat_
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/const_holder_conv.py,sha256=sw3jIOUSvfWUeD8l3rGcUOtC6QuzpMIQm8V3RQAM53Q,4741
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_batch_norm.py,sha256=7GZY7lU3LUUaO5iiccHkUP62PB0QeGAGOZdUSGMkFBY,4450
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_layer_norm.py,sha256=XhiLVcnCc_gF-6mjxbf9C4bYg5YL_GCvDJmcdLkBNAg,4151
+model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_linear.py,sha256=3-OHYPun5Rt7GITqV3ZekJk59tsuY9ZYSpRpxKsNEVA,3450
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/linear_collapsing.py,sha256=CXSMASpc_Zed3BJ2CsER69zKxE6ncFvvKQWDO1JxKYI,5849
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/multi_head_attention_decomposition.py,sha256=VNg-VgzCxSyqy2J3neEPl6U0SPO8UIVU_T47bGhz4FE,38459
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/relu_bound_to_power_of_2.py,sha256=q1a3HieQtaOmWG2WGXp6GHYAvxa3CZ9dJUx9dqMAsS8,5695
@@ -477,9 +478,9 @@ model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tpc_keras.py,sha256=XM6qBLIvzsmdFf-AZq5WOlORK2GXC_X-gulReNxHb9E,6601
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tpc_pytorch.py,sha256=nP05jqvh6uaj30a3W7zEkJfKtqfP0Nz5bobwRqbYrdM,5807
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/__init__.py,sha256=tHTUvsaerSfbe22pU0kIDauPpFD7Pq5EmZytVIDkHz4,717
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tp_model.py,sha256=Ee7M3YVymdv6HYsm7coB8N0dyTOhlAhLdxfSLJXCuoU,15665
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_keras.py,sha256=u8qD1XkHwU4LIoNbmC5mtZd8lZ8gZ4XFihZmoYwAulc,7641
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_pytorch.py,sha256=GCghKkkZOKNTAzwyoZZPid9alGiufNUBzDj2yE7YUSU,6709
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tp_model.py,sha256=b1GAFZMtp0Hf1Ybq8gDLUk90m1HFD00LwtEsFpoN5mY,17240
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_keras.py,sha256=P6gOd5PNMlIp6bcPPfTIX-hTO0AgT9XswrBdvqm-oJ0,8271
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_pytorch.py,sha256=4b3WTV3IDqoqDYx37ba-lxF56K-P5FYyPfIM_TWttQ4,7247
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/target_platform_capabilities.py,sha256=is00rNrDmmirYsyMtMkWz0DwOA92-x7hAJwpd6z1n2E,2806
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/latest/__init__.py,sha256=CXC-HQolSDu7j8V-Xm-SWGCd74gXB3XnAkEhI_TVbIQ,1516
@@ -558,8 +559,8 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=bOc-hFL3gdoSM1Th_S2N_-9JJSlPGpZCTx_QLJHS6lg,3388
 model_compression_toolkit/xquant/pytorch/similarity_functions.py,sha256=CERxq5K8rqaiE-DlwhZBTUd9x69dtYJlkHOPLB54vm8,2354
 model_compression_toolkit/xquant/pytorch/tensorboard_utils.py,sha256=mkoEktLFFHtEKzzFRn_jCnxjhJolK12TZ5AQeDHzUO8,9767
-mct_nightly-2.2.0.20241119.516.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20241119.516.dist-info/METADATA,sha256=-P-wCECFkhsiwVycbvgMYHIWF8OEJLWA_watkr62kFs,26472
-mct_nightly-2.2.0.20241119.516.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
-mct_nightly-2.2.0.20241119.516.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.2.0.20241119.516.dist-info/RECORD,,
+mct_nightly-2.2.0.20241121.524.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20241121.524.dist-info/METADATA,sha256=TpyybcMe4-UNnUKIqO_9NdHs7z0vYVj_7Y7wOs3VaYw,26472
+mct_nightly-2.2.0.20241121.524.dist-info/WHEEL,sha256=bFJAMchF8aTQGUgMZzHJyDDMPTO3ToJ7x23SLJa1SVo,92
+mct_nightly-2.2.0.20241121.524.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.2.0.20241121.524.dist-info/RECORD,,

model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.2.0.20241119.000516"
+__version__ = "2.2.0.20241121.000524"

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_linear.py

@@ -0,0 +1,83 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from torch import nn
+import torch.nn.functional as F
+
+from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
+from model_compression_toolkit.core.common import BaseNode, Graph, BaseSubstitution
+from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
+from model_compression_toolkit.core.pytorch.constants import *
+from model_compression_toolkit.logger import Logger
+
+
+class FunctionalLinear(BaseSubstitution):
+    """
+    Replace functional linear with Linear.
+    """
+
+    def __init__(self):
+        """
+        Matches: functional linear
+        """
+        func_node = NodeOperationMatcher(F.linear)
+        super().__init__(matcher_instance=func_node)
+
+    def substitute(self,
+                   graph: Graph,
+                   func_node: FunctionalNode) -> Graph:
+        """
+        Substitute functional.linear and its inputs with Linear.
+        Args:
+            graph: Graph we apply the substitution on.
+            node: node that match the pattern in the substitution init.
+
+        Returns:
+            Graph after applying the substitution.
+        """
+
+        # Create new node of layer Linear
+        if 1 not in func_node.weights:
+            Logger.critical(f'Weight input missing for node {func_node.name}.')  # pragma: no cover
+        # Extract index of kernel and bias according to tensor_input_allocs if they were input as kwargs. If
+        # they were input as args, use their fixed positions.
+        weight_index = func_node.tensor_input_allocs.index(KERNEL) if KERNEL in func_node.tensor_input_allocs else 1
+        bias_index = func_node.tensor_input_allocs.index(BIAS) if BIAS in func_node.tensor_input_allocs else 2
+        if weight_index not in func_node.weights:
+            Logger.critical(f'Mismatch between tensor_input_allocs and weight index in node {func_node.name}.')  # pragma: no cover
+        weight = func_node.weights[weight_index]
+        bias = func_node.weights.get(bias_index)
+
+        framework_attr = {
+            IN_FEATURES: func_node.input_shape[0][-1],
+            OUT_FEATURES: func_node.output_shape[0][-1],
+            BIAS: bias is not None,
+        }
+
+        weights = {KERNEL: weight} if bias is None else {KERNEL: weight, BIAS: bias}
+
+        new_node = BaseNode(
+            name=func_node.name,
+            framework_attr=framework_attr,
+            input_shape=func_node.input_shape[0],
+            output_shape=func_node.output_shape,
+            weights=weights,
+            layer_class=nn.Linear,
+            has_activation=func_node.has_activation,
+            reuse=func_node.reuse,
+            reuse_group=func_node.reuse_group
+        )
+
+        graph.replace_node(func_node, new_node)
+        return graph

model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -50,6 +50,8 @@ from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.fu
     FunctionalBatchNorm
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.functional_layer_norm import \
     FunctionalLayerNorm
+from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.functional_linear import \
+    FunctionalLinear
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.linear_collapsing import \
     pytorch_linear_collapsing
 from model_compression_toolkit.core.pytorch.graph_substitutions.substitutions.multi_head_attention_decomposition \
@@ -266,6 +268,7 @@ class PytorchImplementation(FrameworkImplementation):
                 FunctionalConvSubstitution(fw_info),
                 FunctionalBatchNorm(),
                 FunctionalLayerNorm(),
+                FunctionalLinear(),
                 RemoveIdentity()]
 
     def get_substitutions_pre_statistics_collection(self,

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

@@ -31,15 +31,20 @@ OPSET_DIMENSION_MANIPULATION_OPS = "DimensionManipulationOps"
 OPSET_MERGE_OPS = "MergeOps"
 OPSET_CONV = "Conv"
 OPSET_FULLY_CONNECTED = "FullyConnected"
+OPSET_BATCH_NORM = "BatchNorm"
 OPSET_ANY_RELU = "AnyReLU"
 OPSET_ADD = "Add"
 OPSET_SUB = "Sub"
 OPSET_MUL = "Mul"
 OPSET_DIV = "Div"
+OPSET_MIN_MAX = "MinMax"
 OPSET_PRELU = "PReLU"
 OPSET_SWISH = "Swish"
 OPSET_SIGMOID = "Sigmoid"
 OPSET_TANH = "Tanh"
+OPSET_GELU = "Gelu"
+OPSET_HARDSIGMOID = "HardSigmoid"
+OPSET_HARDSWISH = "HardSwish"
 
 
 def get_tp_model() -> TargetPlatformModel:
@@ -172,6 +177,11 @@ def generate_tp_model(default_config: OpQuantizationConfig,
     # If the QuantizationConfigOptions contains only one configuration,
     # this configuration will be used for the operation quantization:
     default_configuration_options = tp.QuantizationConfigOptions([default_config])
+    default_config_input16 = default_config.clone_and_edit(supported_input_activation_n_bits=(8, 16))
+    default_config_options_16bit = tp.QuantizationConfigOptions([default_config_input16,
+                                                                 default_config_input16.clone_and_edit(activation_n_bits=16,
+                                                                                                       signedness=Signedness.SIGNED)],
+                                                                base_config=default_config_input16)
 
     # Create a QuantizationConfigOptions for quantizing constants in functional ops.
     # Constant configuration is similar to the default eight bit configuration except for PoT
@@ -212,6 +222,9 @@ def generate_tp_model(default_config: OpQuantizationConfig,
                                                                weights_per_channel_threshold=False))
     qpreserving_const_config_options = tp.QuantizationConfigOptions([qpreserving_const_config])
 
+    mp_cfg_list_16bit = [mp_cfg.clone_and_edit(activation_n_bits=16, signedness=Signedness.SIGNED)
+                         for mp_cfg in mixed_precision_cfg_list]
+
     # 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):
@@ -246,30 +259,37 @@ def generate_tp_model(default_config: OpQuantizationConfig,
         tp.OperatorsSet(OPSET_MERGE_OPS, const_configuration_options_inout16_per_tensor)
 
         # Create Mixed-Precision quantization configuration options from the given list of OpQuantizationConfig objects
-        mixed_precision_configuration_options = tp.QuantizationConfigOptions(mixed_precision_cfg_list,
+        mixed_precision_configuration_options = tp.QuantizationConfigOptions(mixed_precision_cfg_list + mp_cfg_list_16bit,
                                                                              base_config=base_config)
 
         # Define operator sets that use mixed_precision_configuration_options:
         conv = tp.OperatorsSet(OPSET_CONV, mixed_precision_configuration_options)
         fc = tp.OperatorsSet(OPSET_FULLY_CONNECTED, mixed_precision_configuration_options)
 
-        # Define operations sets without quantization configuration
-        # options (useful for creating fusing patterns, for example):
-        any_relu = tp.OperatorsSet(OPSET_ANY_RELU)
+        tp.OperatorsSet(OPSET_BATCH_NORM, default_config_options_16bit)
+
+        # Note: Operations sets without quantization configuration are useful for creating fusing patterns
+        any_relu = tp.OperatorsSet(OPSET_ANY_RELU, default_config_options_16bit)
         add = tp.OperatorsSet(OPSET_ADD, const_configuration_options_inout16)
         sub = tp.OperatorsSet(OPSET_SUB, const_configuration_options_inout16)
         mul = tp.OperatorsSet(OPSET_MUL, const_configuration_options_inout16)
         div = tp.OperatorsSet(OPSET_DIV, const_configuration_options)
-        prelu = tp.OperatorsSet(OPSET_PRELU)
-        swish = tp.OperatorsSet(OPSET_SWISH)
-        sigmoid = tp.OperatorsSet(OPSET_SIGMOID)
-        tanh = tp.OperatorsSet(OPSET_TANH)
+        tp.OperatorsSet(OPSET_MIN_MAX, const_configuration_options_inout16)
+        prelu = tp.OperatorsSet(OPSET_PRELU, default_config_options_16bit)
+        swish = tp.OperatorsSet(OPSET_SWISH, default_config_options_16bit)
+        sigmoid = tp.OperatorsSet(OPSET_SIGMOID, default_config_options_16bit)
+        tanh = tp.OperatorsSet(OPSET_TANH, default_config_options_16bit)
+        gelu = tp.OperatorsSet(OPSET_GELU, default_config_options_16bit)
+        hardsigmoid = tp.OperatorsSet(OPSET_HARDSIGMOID, default_config_options_16bit)
+        hardswish = tp.OperatorsSet(OPSET_HARDSWISH, default_config_options_16bit)
 
         # 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 = tp.OperatorSetConcat(any_relu, swish, prelu, sigmoid, tanh)
-        activations_after_fc_to_fuse = tp.OperatorSetConcat(any_relu, swish, sigmoid)
+        activations_after_conv_to_fuse = tp.OperatorSetConcat(any_relu, swish, prelu, sigmoid,
+                                                              tanh, gelu, hardswish, hardsigmoid)
+        activations_after_fc_to_fuse = tp.OperatorSetConcat(any_relu, swish, sigmoid, tanh, gelu,
+                                                            hardswish, hardsigmoid)
         any_binary = tp.OperatorSetConcat(add, sub, mul, div)
 
         # ------------------- #

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

@@ -26,11 +26,11 @@ 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
+        Conv2DTranspose, Identity, Concatenate, BatchNormalization, Minimum, Maximum
 else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
-        Conv2DTranspose, Identity, Concatenate
+        Conv2DTranspose, Identity, Concatenate, BatchNormalization, Minimum, Maximum
 
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tp_model import get_tp_model
 import model_compression_toolkit as mct
@@ -38,7 +38,7 @@ from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tp
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tp_model import OPSET_NO_QUANTIZATION, \
     OPSET_QUANTIZATION_PRESERVING, OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, OPSET_DIMENSION_MANIPULATION_OPS, \
     OPSET_MERGE_OPS, OPSET_CONV, OPSET_FULLY_CONNECTED, OPSET_ANY_RELU, OPSET_ADD, OPSET_SUB, OPSET_MUL, OPSET_DIV, \
-    OPSET_PRELU, OPSET_SWISH, OPSET_SIGMOID, OPSET_TANH
+    OPSET_PRELU, OPSET_SWISH, OPSET_SIGMOID, OPSET_TANH, OPSET_GELU, OPSET_BATCH_NORM, OPSET_MIN_MAX, OPSET_HARDSIGMOID
 
 tp = mct.target_platform
 
@@ -117,6 +117,7 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
         tp.OperationsSetToLayers(OPSET_FULLY_CONNECTED, [Dense],
                                  attr_mapping={KERNEL_ATTR: DefaultDict(default_value=KERAS_KERNEL),
                                                BIAS_ATTR: DefaultDict(default_value=BIAS)})
+        tp.OperationsSetToLayers(OPSET_BATCH_NORM, [BatchNormalization])
         tp.OperationsSetToLayers(OPSET_ANY_RELU, [tf.nn.relu,
                                                   tf.nn.relu6,
                                                   tf.nn.leaky_relu,
@@ -128,9 +129,13 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
         tp.OperationsSetToLayers(OPSET_SUB, [tf.subtract, Subtract])
         tp.OperationsSetToLayers(OPSET_MUL, [tf.math.multiply, Multiply])
         tp.OperationsSetToLayers(OPSET_DIV, [tf.math.divide, tf.math.truediv])
+        tp.OperationsSetToLayers(OPSET_MIN_MAX, [tf.math.minimum, tf.math.maximum, Minimum, Maximum])
         tp.OperationsSetToLayers(OPSET_PRELU, [PReLU])
         tp.OperationsSetToLayers(OPSET_SWISH, [tf.nn.swish, tp.LayerFilterParams(Activation, activation="swish")])
         tp.OperationsSetToLayers(OPSET_SIGMOID, [tf.nn.sigmoid, tp.LayerFilterParams(Activation, activation="sigmoid")])
         tp.OperationsSetToLayers(OPSET_TANH, [tf.nn.tanh, tp.LayerFilterParams(Activation, activation="tanh")])
+        tp.OperationsSetToLayers(OPSET_GELU, [tf.nn.gelu, tp.LayerFilterParams(Activation, activation="gelu")])
+        tp.OperationsSetToLayers(OPSET_HARDSIGMOID, [tf.keras.activations.hard_sigmoid,
+                                                     tp.LayerFilterParams(Activation, activation="hard_sigmoid")])
 
     return keras_tpc

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

@@ -17,11 +17,13 @@ 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
-from torch.nn import Conv2d, Linear, ConvTranspose2d, MaxPool2d
+    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, LeakyReLU
-from torch.nn.functional import relu, relu6, prelu, silu, hardtanh, hardswish, leaky_relu
+from torch.nn import ReLU, ReLU6, PReLU, SiLU, Sigmoid, Tanh, Hardswish, Hardsigmoid, LeakyReLU, GELU
+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.defaultdict import DefaultDict
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, BIAS_ATTR, PYTORCH_KERNEL, \
@@ -32,7 +34,8 @@ from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tp
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v4.tp_model import OPSET_NO_QUANTIZATION, \
     OPSET_QUANTIZATION_PRESERVING, OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, OPSET_DIMENSION_MANIPULATION_OPS, \
     OPSET_MERGE_OPS, OPSET_CONV, OPSET_FULLY_CONNECTED, OPSET_ANY_RELU, OPSET_ADD, OPSET_SUB, OPSET_MUL, OPSET_DIV, \
-    OPSET_PRELU, OPSET_SWISH, OPSET_SIGMOID, OPSET_TANH
+    OPSET_PRELU, OPSET_SWISH, OPSET_SIGMOID, OPSET_TANH, OPSET_GELU, OPSET_BATCH_NORM, OPSET_MIN_MAX, OPSET_HARDSIGMOID, \
+    OPSET_HARDSWISH
 
 tp = mct.target_platform
 
@@ -95,6 +98,7 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                  attr_mapping=pytorch_linear_attr_mapping)
         tp.OperationsSetToLayers(OPSET_FULLY_CONNECTED, [Linear],
                                  attr_mapping=pytorch_linear_attr_mapping)
+        tp.OperationsSetToLayers(OPSET_BATCH_NORM, [BatchNorm2d])
         tp.OperationsSetToLayers(OPSET_ANY_RELU, [torch.relu,
                                                   ReLU,
                                                   ReLU6,
@@ -109,9 +113,13 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
         tp.OperationsSetToLayers(OPSET_SUB, [operator.sub, sub, subtract])
         tp.OperationsSetToLayers(OPSET_MUL, [operator.mul, mul, multiply])
         tp.OperationsSetToLayers(OPSET_DIV, [operator.truediv, div, divide])
+        tp.OperationsSetToLayers(OPSET_MIN_MAX, [minimum, maximum])
         tp.OperationsSetToLayers(OPSET_PRELU, [PReLU, prelu])
-        tp.OperationsSetToLayers(OPSET_SWISH, [SiLU, silu, Hardswish, hardswish])
-        tp.OperationsSetToLayers(OPSET_SIGMOID, [Sigmoid, sigmoid])
-        tp.OperationsSetToLayers(OPSET_TANH, [Tanh, tanh])
+        tp.OperationsSetToLayers(OPSET_SWISH, [SiLU, silu])
+        tp.OperationsSetToLayers(OPSET_SIGMOID, [Sigmoid, sigmoid, F.sigmoid])
+        tp.OperationsSetToLayers(OPSET_TANH, [Tanh, tanh, F.tanh])
+        tp.OperationsSetToLayers(OPSET_GELU, [GELU, gelu])
+        tp.OperationsSetToLayers(OPSET_HARDSIGMOID, [Hardsigmoid, hardsigmoid])
+        tp.OperationsSetToLayers(OPSET_HARDSWISH, [Hardswish, hardswish])
 
     return pytorch_tpc