[Feature] NMS update

mmcv/tensorrt/tensorrt_utils.py

@@ -32,18 +32,18 @@ def preprocess_onnx(onnx_model):
 
     init_dict = {_.name: _ for _ in initializers}
 
-    nodes_to_remove = []
+    nodes_name_to_remove = set()
 
     def is_node_without_output(name):
         for node_name, node in node_dict.items():
-            if node not in nodes_to_remove:
+            if node_name not in nodes_name_to_remove:
                 if name in node.input:
                     return False
         return True
 
     def mark_nodes_to_remove(name):
         node = node_dict[name]
-        nodes_to_remove.append(node)
+        nodes_name_to_remove.add(name)
         for input_node_name in node.input:
             if is_node_without_output(input_node_name):
                 mark_nodes_to_remove(input_node_name)
@@ -85,13 +85,16 @@ def parse_data(name, typ, default_value=0):
             if len(node_inputs) >= 3:
                 max_output_boxes_per_class = parse_data(
                     node_inputs[2], np.int64, max_output_boxes_per_class)
+                mark_nodes_to_remove(node_inputs[2])
 
             if len(node_inputs) >= 4:
                 iou_threshold = parse_data(node_inputs[3], np.float32,
                                            iou_threshold)
+                mark_nodes_to_remove(node_inputs[3])
 
             if len(node_inputs) >= 5:
                 score_threshold = parse_data(node_inputs[4], np.float32)
+                mark_nodes_to_remove(node_inputs[4])
 
             new_node = onnx.helper.make_node(
                 'NonMaxSuppression',
@@ -110,8 +113,8 @@ def parse_data(name, typ, default_value=0):
             nodes.insert(idx, new_node)
             nodes.remove(node)
 
-    for node in nodes_to_remove:
-        nodes.remove(node)
+    for node_name in nodes_name_to_remove:
+        nodes.remove(node_dict[node_name])
 
     return onnx_model
 

tests/test_tensorrt_preprocess.py

@@ -0,0 +1,74 @@
+import os
+from functools import wraps
+
+import onnx
+import torch
+
+from mmcv.ops import nms
+from mmcv.tensorrt.tensorrt_utils import preprocess_onnx
+
+
+def remove_tmp_file(func):
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        onnx_file = 'tmp.onnx'
+        kwargs['onnx_file'] = onnx_file
+        result = func(*args, **kwargs)
+        if os.path.exists(onnx_file):
+            os.remove(onnx_file)
+        return result
+
+    return wrapper
+
+
+@remove_tmp_file
+def export_nms_module_to_onnx(module, onnx_file):
+    torch_model = module()
+    torch_model.eval()
+
+    input = (torch.rand([100, 4], dtype=torch.float32),
+             torch.rand([100], dtype=torch.float32))
+
+    torch.onnx.export(
+        torch_model,
+        input,
+        onnx_file,
+        opset_version=11,
+        input_names=['boxes', 'scores'],
+        output_names=['output'])
+
+    onnx_model = onnx.load(onnx_file)
+    onnx.checker.check_model(onnx_model)
+    return onnx_model
+
+
+def test_can_handle_nms_with_constant_maxnum():
+
+    class NMS_with_const_maxnum(torch.nn.Module):
+
+        def forward(self, boxes, scores):
+            return nms(boxes, scores, iou_threshold=0.4, max_num=10)
+
+    onnx_model = export_nms_module_to_onnx(NMS_with_const_maxnum)
+    preprocess_onnx_model = preprocess_onnx(onnx_model)
+    for node in preprocess_onnx_model.graph.node:
+        if 'NonMaxSuppression' in node.name:
+            assert len(node.attribute) == 5, 'The NMS must have 5 attributes.'
+
+
+def test_can_handle_nms_with_undefined_maxnum():
+
+    class NMS_with_const_maxnum(torch.nn.Module):
+
+        def forward(self, boxes, scores):
+            return nms(boxes, scores, iou_threshold=0.4)
+
+    onnx_model = export_nms_module_to_onnx(NMS_with_const_maxnum)
+    preprocess_onnx_model = preprocess_onnx(onnx_model)
+    for node in preprocess_onnx_model.graph.node:
+        if 'NonMaxSuppression' in node.name:
+            assert len(node.attribute) == 5, \
+                'The NMS must have 5 attributes.'
+            assert node.attribute[2].i > 0, \
+                'The max_output_boxes_per_class is not defined correctly.'