Adding max pool to Pytorch IMX500 v3 TPC no quantization list

sony · Jun 17, 2024 · f817147 · f817147
1 parent 06fdc8f
commit f817147
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/..._compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3/tpc_pytorch.py b/..._compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3/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, ConvTranspose2d
+from torch.nn import Conv2d, Linear, ConvTranspose2d, MaxPool2d
 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
@@ -83,7 +83,8 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                     argmax,
                                                     gather,
                                                     topk,
-                                                    squeeze])
+                                                    squeeze,
+                                                    MaxPool2d])
 
         tp.OperationsSetToLayers("Conv", [Conv2d, ConvTranspose2d],
                                  attr_mapping=pytorch_linear_attr_mapping)

diff --git a/...pression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tpc_pytorch.py b/...pression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_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, ConvTranspose2d
+from torch.nn import Conv2d, Linear, ConvTranspose2d, MaxPool2d
 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
@@ -82,7 +82,8 @@ def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
                                                     argmax,
                                                     gather,
                                                     topk,
-                                                    squeeze])
+                                                    squeeze,
+                                                    MaxPool2d])
 
         tp.OperationsSetToLayers("Conv", [Conv2d, ConvTranspose2d],
                                  attr_mapping=pytorch_linear_attr_mapping)