PyTorch usage of INMSLayer on TensorRT


I want to include the NMS function from my Pytorch code to TensorRT conversion (via ONNX). The environment or versions are irrelevant because I want to know how to code it.

There is the option of using plugins batchedNMS or efficientNMS inside the python code. But the problem with these plugins is that they only take 4 variables for boxes and return 4 variable nmsed boxes. No indexing is returned so that we can apply nms to other variables tied to the boxes. So it is not a viable option to use.

TensorRT 8.5.1 introduced NMS functionality via INMSLayer, thankfully it has a selected indexes output. But the problem is that I don’t know how I should code to trigger ONNX to have related node so that it’ll trigger the usage of INMSLayer on TensorRT.

For example Pytorch has a single TopK function and TensorRT introduced support for it in 4.something. But Pytorch does not have an NMS function. Should we be using torchvision.ops.batched_nms to trigger INMSLayer usage in TensorRT? Or torchvision.ops.nms?

Is there any correlation chart to pytorch variants of supported functions?

Thanks in advance,

Still need help on this. What should I use in my Pytorch code to trigger INMSLayer node in TensorRT?