Add "identity" to NoQuantization in TPCs

Remove BatchNorm from NoQuantization
sony · Apr 11, 2024 · 883ec83 · 883ec83
1 parent e295990
commit 883ec83
Show file tree

Hide file tree

Showing 13 changed files with 89 additions and 52 deletions.
diff --git a/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tpc_keras.py b/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tpc_keras.py
@@ -27,10 +27,12 @@
     from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
         Conv2DTranspose
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        from keras.src.layers import Identity
 else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
-        Conv2DTranspose
+        Conv2DTranspose, Identity
 
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tp_model import get_tp_model
 import model_compression_toolkit as mct
@@ -61,7 +63,8 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
 
     keras_tpc = tp.TargetPlatformCapabilities(tp_model, name=name, version=TPC_VERSION)
 
-    no_quant_list = [Reshape,
+    no_quant_list = [tf.identity,
+                     Reshape,
                      tf.reshape,
                      Permute,
                      tf.transpose,
@@ -84,6 +87,9 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
                      tf.image.combined_non_max_suppression,
                      tf.compat.v1.shape]
 
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        no_quant_list.append(Identity)
+
     if FOUND_SONY_CUSTOM_LAYERS:
         no_quant_list.append(SSDPostProcess)
 

diff --git a/..._compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tpc_pytorch.py b/..._compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tpc_pytorch.py
@@ -18,7 +18,7 @@
 import torch
 from torch import add, sub, mul, div, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, chunk, unbind, topk, \
     gather, equal, transpose, permute, argmax, squeeze
-from torch.nn import Conv2d, Linear, BatchNorm2d, ConvTranspose2d
+from torch.nn import Conv2d, Linear, ConvTranspose2d, Identity
 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
@@ -64,15 +64,15 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                    BIAS_ATTR: DefaultDict(default_value=BIAS)}
 
     with pytorch_tpc:
-        tp.OperationsSetToLayers("NoQuantization", [Dropout,
+        tp.OperationsSetToLayers("NoQuantization", [Identity,
+                                                    Dropout,
                                                     Flatten,
                                                     dropout,
                                                     flatten,
                                                     split,
                                                     operator.getitem,
                                                     reshape,
                                                     unsqueeze,
-                                                    BatchNorm2d,
                                                     chunk,
                                                     unbind,
                                                     torch.Tensor.size,

diff --git a/...ompression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_lut/tpc_keras.py b/...ompression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_lut/tpc_keras.py
@@ -27,6 +27,8 @@
     from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
         Conv2DTranspose
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        from keras.src.layers import Identity
 else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
@@ -61,7 +63,8 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
 
     keras_tpc = tp.TargetPlatformCapabilities(tp_model, name=name, version=TPC_VERSION)
 
-    no_quant_list = [Reshape,
+    no_quant_list = [tf.identity,
+                     Reshape,
                      tf.reshape,
                      Permute,
                      tf.transpose,
@@ -84,6 +87,9 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
                      tf.image.combined_non_max_suppression,
                      tf.compat.v1.shape]
 
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        no_quant_list.append(Identity)
+
     if FOUND_SONY_CUSTOM_LAYERS:
         no_quant_list.append(SSDPostProcess)
 

diff --git a/...pression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_lut/tpc_pytorch.py b/...pression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_lut/tpc_pytorch.py
@@ -18,7 +18,7 @@
 import torch
 from torch import add, sub, mul, div, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, chunk, unbind, topk, \
     gather, equal, transpose, permute, argmax, squeeze
-from torch.nn import Conv2d, Linear, BatchNorm2d, ConvTranspose2d
+from torch.nn import Conv2d, Linear, ConvTranspose2d, Identity
 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
@@ -64,15 +64,15 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                    BIAS_ATTR: DefaultDict(default_value=BIAS)}
 
     with pytorch_tpc:
-        tp.OperationsSetToLayers("NoQuantization", [Dropout,
+        tp.OperationsSetToLayers("NoQuantization", [Identity,
+                                                    Dropout,
                                                     Flatten,
                                                     dropout,
                                                     flatten,
                                                     split,
                                                     operator.getitem,
                                                     reshape,
                                                     unsqueeze,
-                                                    BatchNorm2d,
                                                     chunk,
                                                     unbind,
                                                     torch.Tensor.size,

diff --git a/...ompression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tpc_keras.py b/...ompression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tpc_keras.py
@@ -27,10 +27,12 @@
     from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
         Conv2DTranspose
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        from keras.src.layers import Identity
 else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
-        Conv2DTranspose
+        Conv2DTranspose, Identity
 
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1_pot.tp_model import get_tp_model
 import model_compression_toolkit as mct
@@ -61,7 +63,8 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
 
     keras_tpc = tp.TargetPlatformCapabilities(tp_model, name=name, version=TPC_VERSION)
 
-    no_quant_list = [Reshape,
+    no_quant_list = [tf.identity,
+                     Reshape,
                      tf.reshape,
                      Permute,
                      tf.transpose,
@@ -84,6 +87,9 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
                      tf.image.combined_non_max_suppression,
                      tf.compat.v1.shape]
 
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        no_quant_list.append(Identity)
+
     if FOUND_SONY_CUSTOM_LAYERS:
         no_quant_list.append(SSDPostProcess)
 

diff --git a/...pression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tpc_pytorch.py b/...pression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tpc_pytorch.py
@@ -18,7 +18,7 @@
 import torch
 from torch import add, sub, mul, div, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, chunk, unbind, topk, \
     gather, equal, transpose, permute, argmax, squeeze
-from torch.nn import Conv2d, Linear, BatchNorm2d, ConvTranspose2d
+from torch.nn import Conv2d, Linear, ConvTranspose2d, Identity
 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
@@ -65,15 +65,15 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                    BIAS_ATTR: DefaultDict(default_value=BIAS)}
 
     with pytorch_tpc:
-        tp.OperationsSetToLayers("NoQuantization", [Dropout,
+        tp.OperationsSetToLayers("NoQuantization", [Identity,
+                                                    Dropout,
                                                     Flatten,
                                                     dropout,
                                                     flatten,
                                                     split,
                                                     operator.getitem,
                                                     reshape,
                                                     unsqueeze,
-                                                    BatchNorm2d,
                                                     chunk,
                                                     unbind,
                                                     torch.Tensor.size,

diff --git a/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tpc_keras.py b/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tpc_keras.py
@@ -27,10 +27,12 @@
     from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
         Conv2DTranspose
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        from keras.src.layers import Identity
 else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
-        Conv2DTranspose
+        Conv2DTranspose, Identity
 
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2.tp_model import get_tp_model
 import model_compression_toolkit as mct
@@ -62,7 +64,8 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
 
     keras_tpc = tp.TargetPlatformCapabilities(tp_model, name=name, version=TPC_VERSION)
 
-    no_quant_list = [Reshape,
+    no_quant_list = [tf.identity,
+                     Reshape,
                      tf.reshape,
                      Permute,
                      tf.transpose,
@@ -85,6 +88,9 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
                      tf.image.combined_non_max_suppression,
                      tf.compat.v1.shape]
 
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        no_quant_list.append(Identity)
+
     if FOUND_SONY_CUSTOM_LAYERS:
         no_quant_list.append(SSDPostProcess)
 

diff --git a/..._compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tpc_pytorch.py b/..._compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tpc_pytorch.py
@@ -18,7 +18,7 @@
 import torch
 from torch import add, sub, mul, div, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, chunk, unbind, topk, \
     gather, equal, transpose, permute, argmax, squeeze
-from torch.nn import Conv2d, Linear, BatchNorm2d, ConvTranspose2d
+from torch.nn import Conv2d, Linear, ConvTranspose2d
 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
@@ -73,7 +73,6 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                     operator.getitem,
                                                     reshape,
                                                     unsqueeze,
-                                                    BatchNorm2d,
                                                     chunk,
                                                     unbind,
                                                     torch.Tensor.size,

diff --git a/...ompression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tpc_keras.py b/...ompression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tpc_keras.py
@@ -27,10 +27,12 @@
     from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
         Conv2DTranspose
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        from keras.src.layers import Identity
 else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
         MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
-        Conv2DTranspose
+        Conv2DTranspose, Identity
 
 from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v2_lut.tp_model import get_tp_model
 import model_compression_toolkit as mct
@@ -61,7 +63,8 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
 
     keras_tpc = tp.TargetPlatformCapabilities(tp_model, name=name, version=TPC_VERSION)
 
-    no_quant_list = [Reshape,
+    no_quant_list = [tf.identity,
+                     Reshape,
                      tf.reshape,
                      Permute,
                      tf.transpose,
@@ -84,6 +87,9 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
                      tf.image.combined_non_max_suppression,
                      tf.compat.v1.shape]
 
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        no_quant_list.append(Identity)
+
     if FOUND_SONY_CUSTOM_LAYERS:
         no_quant_list.append(SSDPostProcess)
 

diff --git a/...pression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tpc_pytorch.py b/...pression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tpc_pytorch.py
@@ -18,7 +18,7 @@
 import torch
 from torch import add, sub, mul, div, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, chunk, unbind, topk, \
     gather, equal, transpose, permute, argmax, squeeze
-from torch.nn import Conv2d, Linear, BatchNorm2d, ConvTranspose2d
+from torch.nn import Conv2d, Linear, ConvTranspose2d
 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
@@ -72,7 +72,6 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                     operator.getitem,
                                                     reshape,
                                                     unsqueeze,
-                                                    BatchNorm2d,
                                                     chunk,
                                                     unbind,
                                                     torch.Tensor.size,

diff --git a/model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tpc_keras.py b/model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tpc_keras.py
@@ -21,9 +21,11 @@
 if version.parse(tf.__version__) >= version.parse("2.13"):
     from keras.src.layers import Conv2D, Dense, Reshape, ZeroPadding2D, AveragePooling2D, Activation, DepthwiseConv2D, \
         MaxPooling2D, ReLU, Add, Softmax, Concatenate, Multiply, Maximum, Minimum, BatchNormalization
+    if version.parse(tf.__version__) >= version.parse("2.12"):
+        from keras.src.layers import Identity
 else:
     from keras.layers import Conv2D, Dense, Reshape, ZeroPadding2D, AveragePooling2D, Activation, DepthwiseConv2D, \
-        MaxPooling2D, ReLU, Add, Softmax, Concatenate, Multiply, Maximum, Minimum, BatchNormalization
+        MaxPooling2D, ReLU, Add, Softmax, Concatenate, Multiply, Maximum, Minimum, BatchNormalization, Identity
 
 from tensorflow.python.keras.layers.core import SlicingOpLambda
 from tensorflow.python.ops.image_ops_impl import ResizeMethod
@@ -62,31 +64,37 @@ def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                               version=TPC_VERSION)
 
     with keras_tpc:
-        tp.OperationsSetToLayers("NoQuantization", [AveragePooling2D,
-                                                    tf.nn.avg_pool2d,
-                                                    Concatenate,
-                                                    tf.concat,
-                                                    MaxPooling2D,
-                                                    Multiply,
-                                                    tf.multiply,
-                                                    Reshape,
-                                                    tf.reshape,
-                                                    tp.LayerFilterParams(tf.image.resize,
-                                                                         method=ResizeMethod.BILINEAR),
-                                                    tf.nn.space_to_depth,
-                                                    ZeroPadding2D,
-                                                    tf.unstack,
-                                                    tf.gather,
-                                                    tf.compat.v1.batch_to_space_nd,
-                                                    tf.space_to_batch_nd,
-                                                    tf.transpose,
-                                                    tf.maximum,
-                                                    Maximum,
-                                                    tf.minimum,
-                                                    Minimum,
-                                                    tf.pad,
-                                                    tf.slice,
-                                                    SlicingOpLambda])
+        no_quant_list = [tf.identity,
+                         AveragePooling2D,
+                         tf.nn.avg_pool2d,
+                         Concatenate,
+                         tf.concat,
+                         MaxPooling2D,
+                         Multiply,
+                         tf.multiply,
+                         Reshape,
+                         tf.reshape,
+                         tp.LayerFilterParams(tf.image.resize,
+                                              method=ResizeMethod.BILINEAR),
+                         tf.nn.space_to_depth,
+                         ZeroPadding2D,
+                         tf.unstack,
+                         tf.gather,
+                         tf.compat.v1.batch_to_space_nd,
+                         tf.space_to_batch_nd,
+                         tf.transpose,
+                         tf.maximum,
+                         Maximum,
+                         tf.minimum,
+                         Minimum,
+                         tf.pad,
+                         tf.slice,
+                         SlicingOpLambda]
+
+        if version.parse(tf.__version__) >= version.parse("2.12"):
+            no_quant_list.append(Identity)
+
+        tp.OperationsSetToLayers("NoQuantization", no_quant_list)
 
         tp.OperationsSetToLayers("FullyConnected", [Dense],
                                  # we provide attributes mapping that maps each layer type in the operations set

diff --git a/..._compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tpc_pytorch.py b/..._compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tpc_pytorch.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import torch
-from torch.nn import AvgPool2d, MaxPool2d
+from torch.nn import AvgPool2d, MaxPool2d, Identity
 from torch.nn.functional import avg_pool2d, max_pool2d, interpolate
 
 from model_compression_toolkit.defaultdict import DefaultDict
@@ -51,7 +51,8 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                 version=TPC_VERSION)
 
     with pytorch_tpc:
-        tp.OperationsSetToLayers("NoQuantization", [AvgPool2d,
+        tp.OperationsSetToLayers("NoQuantization", [Identity,
+                                                    AvgPool2d,
                                                     avg_pool2d,
                                                     torch.cat,
                                                     torch.concat,