CombinedNonMaxSuppression is not supported in ONNX

[jan-golda]

Describe the bug
I am trying to convert MaskRCNN in TensorFlow 2 to ONNX but it is failing due to CombinedNonMaxSuppression op being not supported in ONNX.

Urgency
It is blocking the use of MaskRCNN in ONNX.

System information

• OS Platform and Distribution: Linux Ubuntu 20.04.1
• Tensorflow Version: 2.3.1
• ONNX version: 1.8.0
• Python version: 3.8.5

To Reproduce
Try to convert any model that uses tf.image.combined_non_max_suppression

Expected behaviour
A model should be converted without a fail on CombinedNonMaxSuppression op.

Additional context
This was already reported in #847 for YOLO, and I have tried to apply the workaround from there - replace CombinedNonMaxSuppression with a NonMaxSuppression accompanied by a set of ops that were meant to replace the "Combined" part.

I have tried to get it to work for a few days, but in the case of MaskRCNN, this seems to be more complicated than in the case of YOLO. I had to apply the NonMaxSuppression for each class in each sample from batch separately, then pad it, select the top results for a class, then select the top results for a box, retrieve information about classes and scores, pad it again and finally gather results across the batch.

There is a reason why this op was added to TF, as recreating it from scratch is quite complicated, therefore I would like to ask if you could add the support for it in the ONNX.

Moreover, when I was running the model with partially applied changes I was observing significant performance drop when running it with automatic mixed-precision - in short: replacing the CombinedNonMaxSuppression has a noticeable impact on the original TF model, which is not ideal.

[TomWildenhain-Microsoft]

Hi @jan-golda to convert this op, we can either 1) compose it out of ONNX ops or 2) implement it as a custom op. From your experimenting in TF is seems like you are finding making a compositions of ops to be expensive. We have previously made relatively complicated ops out of compositions of ops with decent performance, but sometimes it can't be done. In this case, I suspect we can get an efficient composition. Are you using any loops in your implementation, or are you using only batched tensor ops + Gather?

[jan-golda]

Hi @TomWildenhain-Microsoft
Well, I was trying to use the NonMaxSuppression, but since it does no support batching nor class-wise suppression I had to iterate over both the batch and the classes.
I've implemented that using two nested map_fn accompanied with a lot of padding/stacking/reshaping. I expect I could just use a for loop for that, no idea what impact this would have on performance.

[guschmue]

Should be possible to add support for CombinedNonMaxSuppression since the onnx nms op supports batch.
We actually unsqueeze the input to get a batch size of 1.

[jan-golda]

@guschmue nice to hear that!

Do you think it would be possible to implement that in the near future?

[TomWildenhain-Microsoft]

Yep, I'm working on it! Hopefully by end of the week.

[TomWildenhain-Microsoft]

What is the dimension of the scores tensor for your model? TF has 2 different behaviors (are you sharing boxes across classes?)

[jan-golda]

Sorry for the late reply!

There are two separate places in the code where this op is used. Below you will find some example shapes for these two places:

place 1:
  boxes:  [4, 1000, 90, 4]
  scores: [4, 1000, 90]
place 2:
  boxes:  [4, 209664, 1, 4]
  scores: [4, 209664, 1]

So the answer is no - I am not sharing the boxes across classes since the third dimension of boxes always equals the third dimension of scores

[TomWildenhain-Microsoft]

Well... just finished implementing it for the other version: #1376. Non-sharing is a little harder since ONNX does share boxes across classes for NMS. I could just make the boxes shared and zero-out the score for all but one class, but that will be a lot of zeros (90 per box, so 90 * 4 * 1000 * 90 = 32 million (probably too large). Place 2 should work fine with the current implementation since there is only 1 class. For place 1, what is the max_outputs and max outputs per class?

[jan-golda]

max_output_size_per_class=1000,
max_total_size=100

[TomWildenhain-Microsoft]

For experimental purposes, can you try testing the performance of the CombinedNonMaxSuppression implementation I've done so far? Making boxes non-shared will be a decent bit harder but have similar perf so it would be nice to know if the perf is sufficient. Just add a slice before your CombinedNonMaxSuppression to cut the dim from 90 to 1 and see how the perf of the conversion to ONNX compares to TF.

If the perf is not good, we may have to use a custom op or try a different implementation approach.

[guschmue]

assume this is resolved.

[PINTO0309]

Pure ONNX Multi-Class NonMaximumSuppression, CombinedNonMaxSuppression.
https://github.com/PINTO0309/yolact_edge_onnx_tensorrt_myriad