Class-agnostic NMS plugin


It would be great if you could include an option to make the BatchedNMS plugin perform NMS filtering in a class-agnostic way. Right now, the plugin only filters based on IoU scores among detections belonging to the same class - what I would like is an option to do this across all detections, regardless of the class.


In my case, the detector has trouble distinguishing between two classes. That’s why it can produce detections only a few pixels apart, with some having the highest score for class A, and others for class B. The class-wise way of NMS filtering will unfortunately then return two such almost identical detection boxes, except that one has class A, the other class B - although they are both representing the same object.

In my dataset I found that this actually significantly hurts the precision scores of affected classes.


The best way would be to add classAgnostic or similar attribute to the plugin. I’ve tried everything I could think of - it’s not possible to achieve what I want with the current plugin (giving all instances the same class means that there is no way to recover the actual class). I would be glad to help with this, but I have very limited experince with C++ and CUDA, and not that much time :(

This behavior is supported by ONNX NMS operator, which also has limited support in TensorRT (but does not support dynamic shapes, which is why I can’t use it).

Please refer to below links related custom plugin implementation and sample:

While IPluginV2 and IPluginV2Ext interfaces are still supported for backward compatibility with TensorRT 5.1 and 6.0.x respectively, however, we recommend that you write new plugins or refactor existing ones to target the IPluginV2DynamicExt or IPluginV2IOExt interfaces instead.


Thanks, but what I am suggesting is a change to the existing plugin

I see that EfficientNMS plug in has the same problem.


The EfficientNMS plugin will support class-agnostic filtering in its coming updates. If you are using BatchedNMS currently, it should be easy to switch to EfficientNMS as both plugins share similar input/output tensor specs.

In the meantime, if the you needs this change urgently (and if you are comfortable with editing and recompiling plugin source code) you can do this change easily by removing the following line in the code and rebuilding the plugin library.

threadClass[tile] == testClass &&

Removing that check in the if statement will skip the class-specific filtering, and the plugin would then effectively operate as class-agnostic instead.

Thank you.


@spolisetty Great news, thank you! I have actually already switched to EfficientNMS plugin, as it has a simpler interface and was a bit faster.

Can I ask when can we expect the new version of the plugin with agnostic NMS to be released?

Currently we are not sure about ETA, we can expect it in coming couple of weeks.

Thank you.

1 Like

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.