box decode tpc

Author: yarden-sony

model_compression_toolkit/target_platform_capabilities/schema/v1.py

@@ -74,6 +74,7 @@ class OperatorSetNames(str, Enum):
     TOPK = "TopK"
     FAKE_QUANT = "FakeQuant"
     COMBINED_NON_MAX_SUPPRESSION = "CombinedNonMaxSuppression"
+    BOX_DECODE = "BoxDecode"
     ZERO_PADDING2D = "ZeroPadding2D"
     CAST = "Cast"
     RESIZE = "Resize"

model_compression_toolkit/target_platform_capabilities/targetplatform2framework/attach2pytorch.py

@@ -32,7 +32,8 @@
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2fw import \
     AttachTpcToFramework
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attribute_filter import Eq
-from sony_custom_layers.pytorch import MulticlassNMS, MulticlassNMSWithIndices, multiclass_nms_with_indices
+from sony_custom_layers.pytorch import MulticlassNMS, MulticlassNMSWithIndices, multiclass_nms_with_indices, \
+    FasterRCNNBoxDecode, multiclass_nms
 
 
 class AttachTpcToPytorch(AttachTpcToFramework):
@@ -98,7 +99,9 @@ def __init__(self):
             OperatorSetNames.L2NORM: [LayerFilterParams(torch.nn.functional.normalize,
                                                         Eq('p', 2) | Eq('p', None))],
             OperatorSetNames.SSD_POST_PROCESS: [],  # no such operator in pytorch
-            OperatorSetNames.COMBINED_NON_MAX_SUPPRESSION: [multiclass_nms_with_indices, MulticlassNMS, MulticlassNMSWithIndices]  # no such operator in pytorch
+            OperatorSetNames.COMBINED_NON_MAX_SUPPRESSION: [multiclass_nms, multiclass_nms_with_indices, MulticlassNMS,
+                                                            MulticlassNMSWithIndices],  # no such operator in pytorch
+            OperatorSetNames.BOX_DECODE: [FasterRCNNBoxDecode]
         }
 
         pytorch_linear_attr_mapping = {KERNEL_ATTR: DefaultDict(default_value=PYTORCH_KERNEL),