Move splitting ops from default to qpreserving configs in TPCv4. (#1304)

sony · Dec 31, 2024 · 8863eec · 8863eec
1 parent a5f1899
commit 8863eec
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 10 deletions.
diff --git a/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tp_model.py b/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tp_model.py
@@ -29,6 +29,7 @@
 OPSET_QUANTIZATION_PRESERVING = "QuantizationPreserving"
 OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS = "DimensionManipulationOpsWithWeights"
 OPSET_DIMENSION_MANIPULATION_OPS = "DimensionManipulationOps"
+OPSET_SPLIT_OPS = "SplitOps"
 OPSET_MERGE_OPS = "MergeOps"
 OPSET_CONV = "Conv"
 OPSET_FULLY_CONNECTED = "FullyConnected"
@@ -186,6 +187,15 @@ def generate_tp_model(default_config: OpQuantizationConfig,
                                                                          signedness=Signedness.SIGNED)]),
                                                                     base_config=default_config_input16)
 
+    qpreseving_config = default_config.clone_and_edit(enable_activation_quantization=False,
+                                                      quantization_preserving=True,
+                                                      supported_input_activation_n_bits=(8, 16))
+    qpreseving_config_options = schema.QuantizationConfigOptions(quantization_configurations=tuple([qpreseving_config,
+                                                                        qpreseving_config.clone_and_edit(
+                                                                            activation_n_bits=16,
+                                                                            signedness=Signedness.SIGNED)]),
+                                                                 base_config=qpreseving_config)
+
     # Create a QuantizationConfigOptions for quantizing constants in functional ops.
     # Constant configuration is similar to the default eight bit configuration except for PoT
     # quantization method for the constant.
@@ -260,6 +270,7 @@ def generate_tp_model(default_config: OpQuantizationConfig,
                                                 quantization_preserving=True,
                                                 supported_input_activation_n_bits=(8, 16))
                                             .clone_and_edit_weight_attribute(enable_weights_quantization=False)))
+    operator_set.append(schema.OperatorsSet(name=OPSET_SPLIT_OPS, qc_options=qpreseving_config_options))
     operator_set.append(schema.OperatorsSet(name=OPSET_MERGE_OPS, qc_options=const_configuration_options_inout16_per_tensor))
 
     # Define operator sets that use mixed_precision_configuration_options:

diff --git a/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_keras.py b/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_keras.py
@@ -39,7 +39,8 @@
 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_GELU, OPSET_BATCH_NORM, OPSET_MIN_MAX, OPSET_HARDSIGMOID
+    OPSET_PRELU, OPSET_SWISH, OPSET_SIGMOID, OPSET_TANH, OPSET_GELU, OPSET_BATCH_NORM, OPSET_MIN_MAX, OPSET_HARDSIGMOID, \
+    OPSET_SPLIT_OPS
 
 tp = mct.target_platform
 
@@ -78,11 +79,7 @@ def generate_keras_tpc(name: str, tp_model: schema.TargetPlatformModel):
                                ZeroPadding2D,
                                Dropout,
                                MaxPooling2D,
-                               tf.split,
-                               tf.cast,
-                               tf.unstack,
-                               tf.__operators__.getitem,
-                               tf.strided_slice]
+                               tf.cast]
     quantization_preserving_list_16bit_input = [Reshape,
                                                 tf.reshape,
                                                 Permute,
@@ -97,6 +94,7 @@ def generate_keras_tpc(name: str, tp_model: schema.TargetPlatformModel):
         tp.OperationsSetToLayers(OPSET_QUANTIZATION_PRESERVING, quantization_preserving)
         tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS, quantization_preserving_list_16bit_input)
         tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, [tf.gather, tf.compat.v1.gather])
+        tp.OperationsSetToLayers(OPSET_SPLIT_OPS,[tf.unstack, tf.split, tf.strided_slice, tf.__operators__.getitem])
         tp.OperationsSetToLayers(OPSET_MERGE_OPS, [tf.stack, tf.concat, Concatenate])
         tp.OperationsSetToLayers(OPSET_CONV,
                                  [Conv2D,

diff --git a/..._compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_pytorch.py b/..._compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_pytorch.py
@@ -36,7 +36,7 @@
     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_GELU, OPSET_BATCH_NORM, OPSET_MIN_MAX, OPSET_HARDSIGMOID, \
-    OPSET_HARDSWISH
+    OPSET_HARDSWISH, OPSET_SPLIT_OPS
 
 tp = mct.target_platform
 
@@ -77,9 +77,6 @@ def generate_pytorch_tpc(name: str, tp_model: schema.TargetPlatformModel):
                                                          topk])
         tp.OperationsSetToLayers(OPSET_QUANTIZATION_PRESERVING, [Dropout,
                                                                  dropout,
-                                                                 split,
-                                                                 chunk,
-                                                                 unbind,
                                                                  MaxPool2d])
         tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS, [Flatten,
                                                                     flatten,
@@ -90,6 +87,7 @@ def generate_pytorch_tpc(name: str, tp_model: schema.TargetPlatformModel):
                                                                     permute,
                                                                     transpose])
         tp.OperationsSetToLayers(OPSET_DIMENSION_MANIPULATION_OPS_WITH_WEIGHTS, [gather, torch.Tensor.expand])
+        tp.OperationsSetToLayers(OPSET_SPLIT_OPS,[split, chunk, unbind])
         tp.OperationsSetToLayers(OPSET_MERGE_OPS,
                                  [torch.stack, torch.cat, torch.concat, torch.concatenate])