[Relay, TF] Support converting TF combined_nms using Relay all_class_nms

[masahi]

A follow up to #7796 adding support for converting TF combined nms using all_class_nms.
https://www.tensorflow.org/api_docs/python/tf/image/combined_non_max_suppression

TF and ONNX require different output format from combined NMS. To support TF, we need to return outputs differently. Coming up with the right output format for TF turned out to be non trivial, but I found the following worked quite well:

• indices of size (batch_size, num_boxes * num_class, 2)
• scores of size (batch_size, num_boxes * num_class)
• num_detections of size (batch_size,)

indices and scores are outputs from the NMS loop, where the selected box indices/scores come first. They are sorted within a class, but not across classes. Since TF combined NMS returns max_total_boxes out of num_boxes * num_class, the frontend code needs to do topk.

One complication is even if num_boxes * num_class < max_total_size, TF returns max_total_size boxes with zero padding. Since topi/relay topk op would not work when the size of the axis to select topk elements from is less then k, we need to conditionally do topk depending on the value of num_boxes. Since things become complicated when num_boxes can be dynamic, I use all_class_nms code path only when num_boxes is static (which should be no problem in practice). Moreover, all_class_nms does not support q != 1 case. So if q != 1 or the number of boxes is dynamic, we still use the old converter.

In addition to expected speed up from parallelizing the NMS loop across different classes, this change enables running TF detection models with graph runtime, since the output from TF combined nms is always of size max_total_size and we no longer need to do dynamic slicing.

please review @trevor-m@anijain2305

[masahi]

Since I want to use this function to rewrite ONNX models, it is defined at the top level and visible from outside.

[trevor-m]

Thanks @masahi, this is great. Some small comments.

[trevor-m]

Thanks @masahi, it looks good to me now.

Initial results for TF SSD Mobilenet v1 on T4 GPU are giving me 4.56 ms. It is not quite an apples to apples comparison with the previous implementation though, so I will come back with more data soon.

[masahi]

Interesting. I'd be curious to know why it is not going faster than the previous implementation.

If you can dump the input boxes and scores to TF combined nms, I can look at the performance difference via a standalone NMS workload. I can do that very easily for PT but no idea for TF.

[trevor-m]

The input shapes are boxes=(1, 1917, 1, 4) and scores=(1, 1917, 91).

Here are some more numbers for TF SSD Mobilenet v1.

Implementation	Thrust	Latency (ms)
Previous	Yes	3.41
Previous	No	3.85
New	Yes	4.56
New	No	2.98

It looks like without thrust this implementation is much faster!

Even if we are unable to improve the thrust performance in this PR, I still prefer this new implementation because of the static shapes. The performance is still much better than TF anyway which is ~6ms.

[masahi]

lol thrust is the bottleneck now? Note that our TVM-native sorting performance is quite good thanks to #7611, and it works equally well on dynamic inputs, so thrust is no longer a requirement for good performance. Also the new implementation does not call get_valid_count, which uses thrust for exclusive scan on relatively large input (total number of boxes, which in our case num_boxes * num_class).

I see that the number of boxes is quite small in your model, I believe the new implementation would be much faster when the number of boxes is large. Do you have other models, preferably the one with more input boxes?

I've been testing tf2 models from https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md, their ssd mv2 model (FPNLite 320x320) has 12480 boxes, and ssd resnet 50 v1 has more than 50000, and efficient det2 has 110484 boxes.

[masahi]

@trevor-m Also it is worth trying out the graph runtime for the new implementation, since there is no dynamic shape after NMS.

Thrust numbers actually make sense given that now we are running sort on (batch * class, num_boxes) along the second axis. Since batched sort with thrust is implemented via two calls to stable sort, when the axis to sort is small the overhead from the two calls become relatively big. And for small sort, TVM's sort should be faster than Thrust.

In contrast, the previous implementation does sort on (batch, num_boxes * class) along the second axis, so Thrust sort is fast.

[tqchen]

@trevor-m please help manage this PR, you should get merge permission once the ASF account linking is setup

[masahi]

@trevor-m you should be able to merge now.

[trevor-m]

@trevor-m Also it is worth trying out the graph runtime for the new implementation, since there is no dynamic shape after NMS.

Thrust numbers actually make sense given that now we are running sort on (batch * class, num_boxes) along the second axis. Since batched sort with thrust is implemented via two calls to stable sort, when the axis to sort is small the overhead from the two calls become relatively big. And for small sort, TVM's sort should be faster than Thrust.

In contrast, the previous implementation does sort on (batch, num_boxes * class) along the second axis, so Thrust sort is fast.

I see, that makes sense.

I was using the models from https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1_detection_zoo.md
I think these ones still cannot use graph runtime because there is a loop at the beginning.

[trevor-m]

Thanks @masahi !