Add "identity" to NoQuantization in IMX500.v2 TPC and remove BatchNorm (

#1040) In IMX500.v2 TPC: * Add "identity" to NoQuantization * Remove BatchNorm from Pytorch TPC Note: TF2.11 will fail for that TPC
sony · Apr 15, 2024 · 932a8cb · 932a8cb
1 parent bfb89bd
commit 932a8cb
Show file tree

Hide file tree

Showing 11 changed files with 25 additions and 21 deletions.
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
@@ -26,11 +26,11 @@
 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
+        Conv2DTranspose, 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 +62,9 @@ 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 = [Identity,
+                     tf.identity,
+                     Reshape,
                      tf.reshape,
                      Permute,
                      tf.transpose,

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,8 +18,8 @@
 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 Dropout, Flatten, Hardtanh
+from torch.nn import Conv2d, Linear, ConvTranspose2d
+from torch.nn import Dropout, Flatten, Hardtanh, Identity
 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/...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
@@ -26,11 +26,11 @@
 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
+        Conv2DTranspose, 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 +61,9 @@ 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 = [Identity,
+                     tf.identity,
+                     Reshape,
                      tf.reshape,
                      Permute,
                      tf.transpose,

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,8 +18,8 @@
 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 Dropout, Flatten, Hardtanh
+from torch.nn import Conv2d, Linear, ConvTranspose2d
+from torch.nn import Dropout, Flatten, Hardtanh, Identity
 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/tests/keras_tests/feature_networks_tests/feature_networks/network_editor/node_filter_test.py b/tests/keras_tests/feature_networks_tests/feature_networks/network_editor/node_filter_test.py
@@ -24,7 +24,7 @@
 from model_compression_toolkit.core.common.quantization.quantization_params_fn_selection import \
     get_weights_quantization_params_fn, get_activation_quantization_params_fn
 from model_compression_toolkit.core.keras.constants import KERNEL
-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.latest import generate_keras_tpc
 from tests.common_tests.helpers.generate_test_tp_model import generate_test_tp_model
 from tests.keras_tests.feature_networks_tests.base_keras_feature_test import BaseKerasFeatureNetworkTest
 from tests.keras_tests.utils import get_layers_from_model_by_type

diff --git a/tests/pytorch_tests/exporter_tests/base_pytorch_export_test.py b/tests/pytorch_tests/exporter_tests/base_pytorch_export_test.py
@@ -24,7 +24,7 @@
 from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_onnx_pytorch_exporter import \
     DEFAULT_ONNX_OPSET_VERSION
 from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
-from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import \
+from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import \
     generate_pytorch_tpc
 from tests.common_tests.helpers.generate_test_tp_model import generate_test_tp_model
 

diff --git a/...ytorch_tests/exporter_tests/custom_ops_tests/test_export_lut_symmetric_onnx_quantizers.py b/...ytorch_tests/exporter_tests/custom_ops_tests/test_export_lut_symmetric_onnx_quantizers.py
@@ -16,7 +16,7 @@
 import numpy as np
 import torch
 from mct_quantizers import QuantizationMethod
-from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import \
+from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import \
     generate_pytorch_tpc
 from tests.common_tests.helpers.generate_test_tp_model import generate_test_tp_model
 from tests.pytorch_tests.exporter_tests.base_pytorch_onnx_export_test import BasePytorchONNXCustomOpsExportTest

diff --git a/tests/pytorch_tests/exporter_tests/custom_ops_tests/test_export_pot_onnx_quantizers.py b/tests/pytorch_tests/exporter_tests/custom_ops_tests/test_export_pot_onnx_quantizers.py
@@ -25,7 +25,7 @@
 from model_compression_toolkit.constants import FOUND_ONNX, FOUND_ONNXRUNTIME
 from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_onnx_pytorch_exporter import \
     DEFAULT_ONNX_OPSET_VERSION
-from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import \
+from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import \
     generate_pytorch_tpc
 from tests.common_tests.helpers.generate_test_tp_model import generate_test_tp_model
 from tests.pytorch_tests.exporter_tests.base_pytorch_onnx_export_test import BasePytorchONNXCustomOpsExportTest

diff --git a/tests/pytorch_tests/exporter_tests/custom_ops_tests/test_export_symmetric_onnx_quantizers.py b/tests/pytorch_tests/exporter_tests/custom_ops_tests/test_export_symmetric_onnx_quantizers.py
@@ -26,7 +26,7 @@
 from model_compression_toolkit.constants import FOUND_ONNX, FOUND_ONNXRUNTIME
 from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_onnx_pytorch_exporter import \
     DEFAULT_ONNX_OPSET_VERSION
-from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import \
+from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import \
     generate_pytorch_tpc
 from tests.common_tests.helpers.generate_test_tp_model import generate_test_tp_model
 from tests.pytorch_tests.exporter_tests.base_pytorch_onnx_export_test import BasePytorchONNXCustomOpsExportTest

diff --git a/tests/pytorch_tests/exporter_tests/custom_ops_tests/test_export_uniform_onnx_quantizers.py b/tests/pytorch_tests/exporter_tests/custom_ops_tests/test_export_uniform_onnx_quantizers.py
@@ -17,7 +17,7 @@
 import torch
 
 from mct_quantizers import QuantizationMethod
-from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import \
+from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import \
     generate_pytorch_tpc
 from tests.common_tests.helpers.generate_test_tp_model import generate_test_tp_model
 from tests.pytorch_tests.exporter_tests.base_pytorch_onnx_export_test import BasePytorchONNXCustomOpsExportTest

diff --git a/tests/pytorch_tests/exporter_tests/test_exporting_qat_models.py b/tests/pytorch_tests/exporter_tests/test_exporting_qat_models.py
@@ -24,7 +24,7 @@
 from model_compression_toolkit.constants import FOUND_ONNX, FOUND_ONNXRUNTIME
 from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
-from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.v1.tpc_pytorch import \
+from model_compression_toolkit.target_platform_capabilities.tpc_models.imx500_tpc.latest import \
     generate_pytorch_tpc
 from tests.common_tests.helpers.generate_test_tp_model import generate_test_tp_model
 from tests.pytorch_tests.model_tests.feature_models.qat_test import dummy_train