Batched NMS generates onnx IF operator, which causes onnx conversion problem

🐛 Describe the bug

Background

I'd like to convert Mask R-CNN model into RyzenAI model and run on Ryzen NPU.
When I convert torchvision Mask R-CNN model into ONNX, it contains IF operator.
RyzenAI doesn't support IF operator with NPU inference, so I tried to remove it, but I couldn't found the solution to fix.

Issue

Could anyone please check and answer?

• Is this behavior expected?
• To convert Mask R-CNN into ONNX model without using IF operator, what should I do?

Sample code

import torch
import torchvision

model = torchvision.models.detection.maskrcnn_resnet50_fpn_v2(weights=None, progress=True)
model.load_state_dict(torch.load('./maskrcnn_resnet50_fpn_v2_coco-73cbd019.pth')) # already downloaded from https://download.pytorch.org/models/maskrcnn_resnet50_fpn_v2_coco-73cbd019.pth
model.eval()

x = torch.rand(1, 3, 480, 640)
predictions = model(x)

torch.onnx.export(model,
                  x, 
                  "mask_rcnn_torchvision.onnx",
                  opset_version = 11,
                  input_names=['input'],
                  output_names = ['boxes', 'labels', 'scores', 'masks'],
                  dynamic_axes = {'input': {2: 'height', 3: 'width'}, 'boxes': { 1:'num'}, 'labels':{ 1: 'num'}, 'scores': {1: 'num'},'masks' : {1:'num'}})

Results

mask_rcnn_torchvision.onnx contains IF operator. While this operator is in rpn block, the name is not /rpn/foobar but onnx::If_xxxx.

sub_graph1 (ELSE subgraph):

sub_graph (THEN subgraph):

Expected behavior

mask_rcnn_torchvision.onnx doesn't contain IF operator.

My investigation

I traced inside torchvision Mask R-CNN structure and found that it may be caused by torch.ops.batched_nms()

https://github.com/pytorch/vision/blob/0d68c7df8640abff43355afd57c494cf5d74f4a9/torchvision/ops/boxes.py#L44-L113

I think IF operator is generated by this code:
https://github.com/pytorch/vision/blob/0d68c7df8640abff43355afd57c494cf5d74f4a9/torchvision/ops/boxes.py#L89-L90

However, commented out this code doesn't work, using _batched_nms_vanilla() instead of _batched_nms_coordinate_trick() also doesn't work, too.

Versions

(ryzen-ai) PS C:\work\mask_vision> python .\collect_env.py
Collecting environment information...
PyTorch version: 2.3.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Enterprise Evaluation (10.0.22631 64 bit)
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.31.2
Libc version: N/A

Python version: 3.10.16 | packaged by conda-forge | (main, Dec  5 2024, 14:07:43) [MSC v.1942 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22631-SP0
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Name: AMD Ryzen 7 8845HS
[Detail information is REDACTED because of NDA]

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] onnx==1.16.1
[pip3] onnxruntime==1.17.0
[pip3] onnxruntime_extensions==0.13.0
[pip3] onnxruntime-vitisai==1.19.0
[pip3] onnxscript==0.1.0.dev20250123
[pip3] onnxsim==0.4.36
[pip3] torch==2.3.1+cpu
[pip3] torch-geometric==2.6.1
[pip3] torchvision==0.18.1+cpu
[pip3] vai-q-onnx==1.19.0
[conda] mkl                       2021.4.0                 pypi_0    pypi
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] torch                     2.3.1+cpu                pypi_0    pypi
[conda] torch-geometric           2.6.1                    pypi_0    pypi
[conda] torchvision               0.18.1+cpu               pypi_0    pypi
(ryzen-ai) PS C:\work\mask_vision>

NOTE: It is difficult to update into newest version because of compatibility. It seems that batched_nms is not changed, so anyone may be able to reproduce it in newest version.

Start by reproducing the supplied Mask R-CNN export with the versions and opset shown, then inspect torchvision/ops/boxes.py around batched_nms and the generated ONNX IF node. Compare the coordinate-trick and vanilla paths and determine whether the behavior is expected; done means providing a verified export path without IF or documenting why that is not supported.