[Feature] Added agnostic_nms onnx converter for tensorrt 8

projects/easydeploy/model/model.py

@@ -17,7 +17,7 @@
 from ..backbone import DeployFocus, GConvFocus, NcnnFocus
 from ..bbox_code import (rtmdet_bbox_decoder, yolov5_bbox_decoder,
                          yolox_bbox_decoder)
-from ..nms import batched_nms, efficient_nms, onnx_nms
+from ..nms import agnostic_nms, batched_nms, efficient_nms, onnx_nms
 from .backend import MMYOLOBackend
 
 
@@ -38,6 +38,7 @@ def __init__(self,
             self.with_postprocess = True
             self.__init_sub_attributes()
             self.detector_type = type(self.baseHead)
+            self.agnostic_nms = postprocess_cfg.get('agnostic_nms', False)
             self.pre_top_k = postprocess_cfg.get('pre_top_k', 1000)
             self.keep_top_k = postprocess_cfg.get('keep_top_k', 100)
             self.iou_threshold = postprocess_cfg.get('iou_threshold', 0.65)
@@ -150,7 +151,10 @@ def select_nms(self):
         if self.backend in (MMYOLOBackend.ONNXRUNTIME, MMYOLOBackend.OPENVINO):
             nms_func = onnx_nms
         elif self.backend == MMYOLOBackend.TENSORRT8:
-            nms_func = efficient_nms
+            if self.agnostic_nms:
+                nms_func = agnostic_nms
+            else:
+                nms_func = efficient_nms
         elif self.backend == MMYOLOBackend.TENSORRT7:
             nms_func = batched_nms
         else:

projects/easydeploy/nms/__init__.py

@@ -1,5 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .ort_nms import onnx_nms
-from .trt_nms import batched_nms, efficient_nms
+from .trt_nms import agnostic_nms, batched_nms, efficient_nms
 
-__all__ = ['efficient_nms', 'batched_nms', 'onnx_nms']
+__all__ = ['efficient_nms', 'batched_nms', 'agnostic_nms', 'onnx_nms']

projects/easydeploy/nms/trt_nms.py

@@ -216,6 +216,99 @@ def _batched_nms(
     return num_det, det_boxes, det_scores, det_classes
 
 
+class TRTAgnosticNMSop(torch.autograd.Function):
+
+    @staticmethod
+    def forward(
+        ctx,
+        boxes: Tensor,
+        scores: Tensor,
+        background_class: int = -1,
+        box_coding: int = 0,
+        iou_threshold: float = 0.45,
+        max_output_boxes: int = 100,
+        plugin_version: str = '1',
+        score_activation: int = 0,
+        score_threshold: float = 0.25,
+    ):
+        batch_size, _, num_classes = scores.shape
+        num_det = torch.randint(
+            0, max_output_boxes, (batch_size, 1), dtype=torch.int32)
+        det_boxes = torch.randn(batch_size, max_output_boxes, 4)
+        det_scores = torch.randn(batch_size, max_output_boxes)
+        det_classes = torch.randint(
+            0, num_classes, (batch_size, max_output_boxes), dtype=torch.int32)
+        return num_det, det_boxes, det_scores, det_classes
+
+    @staticmethod
+    def symbolic(g,
+                 boxes: Tensor,
+                 scores: Tensor,
+                 background_class: int = -1,
+                 box_coding: int = 0,
+                 iou_threshold: float = 0.45,
+                 max_output_boxes: int = 100,
+                 plugin_version: str = '1',
+                 score_activation: int = 0,
+                 score_threshold: float = 0.25):
+        out = g.op(
+            'TRT::EfficientNMS_TRT',
+            boxes,
+            scores,
+            background_class_i=background_class,
+            box_coding_i=box_coding,
+            iou_threshold_f=iou_threshold,
+            max_output_boxes_i=max_output_boxes,
+            plugin_version_s=plugin_version,
+            score_activation_i=score_activation,
+            score_threshold_f=score_threshold,
+            outputs=4,
+            class_agnostic_i=1)
+        num_det, det_boxes, det_scores, det_classes = out
+        return num_det, det_boxes, det_scores, det_classes
+
+
+def _agnostic_nms(
+    boxes: Tensor,
+    scores: Tensor,
+    max_output_boxes_per_class: int = 1000,
+    iou_threshold: float = 0.5,
+    score_threshold: float = 0.05,
+    pre_top_k: int = -1,
+    keep_top_k: int = 100,
+    box_coding: int = 0,
+):
+    """Wrapper for `efficient_nms` and `class_agnostic` with TensorRT.
+    Args:
+        boxes (Tensor): The bounding boxes of shape [N, num_boxes, 4].
+        scores (Tensor): The detection scores of shape
+            [N, num_boxes, num_classes].
+        max_output_boxes_per_class (int): Maximum number of output
+            boxes per class of nms. Defaults to 1000.
+        iou_threshold (float): IOU threshold of nms. Defaults to 0.5.
+        score_threshold (float): score threshold of nms.
+            Defaults to 0.05.
+        pre_top_k (int): Number of top K boxes to keep before nms.
+            Defaults to -1.
+        keep_top_k (int): Number of top K boxes to keep after nms.
+            Defaults to -1.
+        box_coding (int): Bounding boxes format for nms.
+            Defaults to 0 means [x1, y1 ,x2, y2].
+            Set to 1 means [x, y, w, h].
+    Returns:
+        tuple[Tensor, Tensor, Tensor, Tensor]:
+        (num_det, det_boxes, det_scores, det_classes),
+        `num_det` of shape [N, 1]
+        `det_boxes` of shape [N, num_det, 4]
+        `det_scores` of shape [N, num_det]
+        `det_classes` of shape [N, num_det]
+    """
+    num_det, det_boxes, det_scores, det_classes = TRTAgnosticNMSop.apply(
+        boxes, scores, -1, box_coding, iou_threshold, keep_top_k, '1', 0,
+        score_threshold)
+    return num_det, det_boxes, det_scores, det_classes
+
+
 def efficient_nms(*args, **kwargs):
     """Wrapper function for `_efficient_nms`."""
     return _efficient_nms(*args, **kwargs)
@@ -224,3 +317,8 @@ def efficient_nms(*args, **kwargs):
 def batched_nms(*args, **kwargs):
     """Wrapper function for `_batched_nms`."""
     return _batched_nms(*args, **kwargs)
+
+
+def agnostic_nms(*args, **kwargs):
+    """Wrapper function for `_agnostic_nms`."""
+    return _agnostic_nms(*args, **kwargs)

projects/easydeploy/tools/export_onnx.py

@@ -51,6 +51,12 @@ def parse_args():
         type=str,
         default='onnxruntime',
         help='Backend for export onnx')
+    parser.add_argument(
+        '--agnostic-nms',
+        action='store_true',
+        help='Switch NMS algorithm to agnostic_nms. ' +
+        'This only works with backend mode with TENSORRT8, ' +
+        'and with TensorRT 8.6 runtime or over.')
     parser.add_argument(
         '--pre-topk',
         type=int,
@@ -94,11 +100,17 @@ def main():
         args.model_only = True
         print_log(f'Can not export postprocess for {args.backend.lower()}.\n'
                   f'Set "args.model_only=True" default.')
+    if args.agnostic_nms and backend != MMYOLOBackend.TENSORRT8:
+        print_log(f'AgnosticNMS only supports TENSORRT8 backend.\n'
+                  'Change your backend from current '
+                  f'{args.backend.lower()} to TENSORRT8.')
+        sys.exit(0)
     if args.model_only:
         postprocess_cfg = None
         output_names = None
     else:
         postprocess_cfg = ConfigDict(
+            agnostic_nms=args.agnostic_nms,
             pre_top_k=args.pre_topk,
             keep_top_k=args.keep_topk,
             iou_threshold=args.iou_threshold,