You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/06/02 10:14:35 UTC

[GitHub] [tvm] masahi opened a new pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

masahi opened a new pull request #8174:
URL: https://github.com/apache/tvm/pull/8174


   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi edited a comment on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-853506007






-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644337411



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,89 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    assert isinstance(batch_size, int), "dynamic batch size not supported yet."
+    tile_batch_reps = _op.const([batch_size, 1])
+    box_range_2d = _op.tile(box_range, tile_batch_reps)
+    valid_mask = _op.cast(
+        _op.less(box_range_2d, _op.expand_dims(num_detections, axis=1)), "float32"
+    )
+
+    def select_topk(do_zero_pad):
+        def true_branch():
+            arange = _op.arange(
+                _op.const(0, dtype="int64"),
+                _op.const(max_output_boxes_per_batch, dtype="int64"),
+                dtype="int64",
+            )
+            pad = _op.full(
+                _op.const(0, dtype="int64"), (max_total_size - max_output_boxes_per_batch,)
+            )
+            topk_indices = _op.tile(_op.concatenate([arange, pad], 0), tile_batch_reps)
+            nmsed_scores = _op.gather(selected_scores, 1, topk_indices)
+            nmsed_scores = nmsed_scores * valid_mask
+            return nmsed_scores, topk_indices
+
+        def false_branch():
+            if isinstance(max_output_boxes_per_class, int):
+                # Do topk on smaller input if possible
+                # TODO(masahi): use axes argument in strided slice when it becomes available
+                slice_mx = _op.const([-1, max_output_boxes_per_class * num_class], dtype="int64")
+                selected_scores_slice = _op.strided_slice(
+                    selected_scores, begin=_op.const([0, 0], dtype="int64"), end=slice_mx

Review comment:
       Thanks for fixing!




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-853506007


   Interesting. I'd be curious to know what 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 test. I can do that very easily for PT but no idea for TF.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854502294


   @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. 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r643921186



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,92 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(

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




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] tqchen commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
tqchen commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854708338


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


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-853506007






-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi edited a comment on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-853506007


   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 test. I can do that very easily for PT but no idea for TF.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi edited a comment on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854239897


   lol thrust is the bottleneck now? Note that our TVM-native sorting performance is quite good thanks to https://github.com/apache/tvm/pull/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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi edited a comment on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854239897


   lol thrust is the bottleneck now? Note that our sorting performance improved thanks to https://github.com/apache/tvm/pull/7611, thrust is no longer a requirement for good performance.
   
   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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644172454



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,92 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    tile_batch_reps = (
+        _op.concatenate([batch_size, 1], axis=0)

Review comment:
       I think if batch_size is an instance of relay.Any() we cannot use batch_size directly here, we need to get the symbolic batch size tensor with `_op.shape_of`.
   
   Could we add a test which has a dynamic batch dimension? If it's too difficult, we can defer to later, since the models I have looked at don't use dynamic batch size for NMS.

##########
File path: python/tvm/topi/transform.py
##########
@@ -483,7 +483,7 @@ def gather(data, axis, indices):
     return cpp.gather(data, axis, indices)
 
 
-def gather_nd(a, indices):
+def gather_nd(a, indices, batch_dims=0):

Review comment:
       Are the gather_nd changes used in this PR?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-855112709


   Thanks @masahi !


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi edited a comment on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854239897


   lol thrust is the bottleneck now? Note that our TVM-native sorting performance is quite good thanks to https://github.com/apache/tvm/pull/7611, 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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644331266



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,89 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    assert isinstance(batch_size, int), "dynamic batch size not supported yet."
+    tile_batch_reps = _op.const([batch_size, 1])
+    box_range_2d = _op.tile(box_range, tile_batch_reps)
+    valid_mask = _op.cast(
+        _op.less(box_range_2d, _op.expand_dims(num_detections, axis=1)), "float32"
+    )
+
+    def select_topk(do_zero_pad):
+        def true_branch():
+            arange = _op.arange(
+                _op.const(0, dtype="int64"),
+                _op.const(max_output_boxes_per_batch, dtype="int64"),
+                dtype="int64",
+            )
+            pad = _op.full(
+                _op.const(0, dtype="int64"), (max_total_size - max_output_boxes_per_batch,)
+            )
+            topk_indices = _op.tile(_op.concatenate([arange, pad], 0), tile_batch_reps)
+            nmsed_scores = _op.gather(selected_scores, 1, topk_indices)
+            nmsed_scores = nmsed_scores * valid_mask
+            return nmsed_scores, topk_indices
+
+        def false_branch():
+            if isinstance(max_output_boxes_per_class, int):
+                # Do topk on smaller input if possible
+                # TODO(masahi): use axes argument in strided slice when it becomes available
+                slice_mx = _op.const([-1, max_output_boxes_per_class * num_class], dtype="int64")
+                selected_scores_slice = _op.strided_slice(
+                    selected_scores, begin=_op.const([0, 0], dtype="int64"), end=slice_mx

Review comment:
       I see, I'd rather use `axes` argument introduced in https://github.com/apache/tvm/pull/8165 after it is merged. This saves us from specifying the begin and end for the batch dim. I was using this in my dev branch, and realized that `axes` argument is not in `main` yet, so added this fishy workaround. And tests in `test_forward.py` have `max_output_boxes_per_class` as a variable, so we don't hit this code path.

##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,89 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    assert isinstance(batch_size, int), "dynamic batch size not supported yet."
+    tile_batch_reps = _op.const([batch_size, 1])
+    box_range_2d = _op.tile(box_range, tile_batch_reps)
+    valid_mask = _op.cast(
+        _op.less(box_range_2d, _op.expand_dims(num_detections, axis=1)), "float32"
+    )
+
+    def select_topk(do_zero_pad):
+        def true_branch():
+            arange = _op.arange(
+                _op.const(0, dtype="int64"),
+                _op.const(max_output_boxes_per_batch, dtype="int64"),
+                dtype="int64",
+            )
+            pad = _op.full(
+                _op.const(0, dtype="int64"), (max_total_size - max_output_boxes_per_batch,)
+            )
+            topk_indices = _op.tile(_op.concatenate([arange, pad], 0), tile_batch_reps)
+            nmsed_scores = _op.gather(selected_scores, 1, topk_indices)
+            nmsed_scores = nmsed_scores * valid_mask
+            return nmsed_scores, topk_indices
+
+        def false_branch():
+            if isinstance(max_output_boxes_per_class, int):
+                # Do topk on smaller input if possible
+                # TODO(masahi): use axes argument in strided slice when it becomes available
+                slice_mx = _op.const([-1, max_output_boxes_per_class * num_class], dtype="int64")
+                selected_scores_slice = _op.strided_slice(
+                    selected_scores, begin=_op.const([0, 0], dtype="int64"), end=slice_mx

Review comment:
       should be fixed

##########
File path: python/tvm/topi/cuda/nms.py
##########
@@ -1012,39 +1079,69 @@ def all_class_non_max_suppression(
     score_threshold : float or tvm.te.Tensor, optional
         Score threshold to filter out low score boxes early
 
+    output_format : str, optional
+        "onnx" or "tensorflow", see below
+
     Returns
     -------
-    out : [tvm.te.Tensor, tvm.te.Tensor]
-        The output is two tensors, the first is `indices` of size
+    out : list of tvm.te.Tensor
+        If `output_format` is "onnx", the output is two tensors. The first is `indices` of size
         `(batch_size * num_class* num_boxes , 3)` and the second is a scalar tensor
         `num_total_detection` of shape `(1,)` representing the total number of selected
-        boxes. Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
+        boxes. The three values in `indices` encode batch, class, and box indices.
+        Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
         first, in descending of scores, followed by boxes from batch 0, class 1 etc. Out of
         `batch_size * num_class* num_boxes` rows of indices, only the first `num_total_detection`
         rows are valid.
+
+        If `output_format` is "tensorflow", the output is three tensors, the first
+        is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the second is `scores` of
+        size `(batch_size, num_class * num_boxes)`, and the third is `num_total_detection` of size
+        `(batch_size,)` representing the total number of selected boxes per batch. The two values
+        in `indices` encode class and box indices. Of num_class * num_boxes boxes in `indices` at
+        batch b, only the first `num_total_detection[b]` entries are valid. The second axis of
+        `indices` and `scores` are sorted within each class by box scores, but not across classes.
+        So the box indices and scores for the class 0 come first in a sorted order, followed by
+        the class 1 etc.
     """
     batch, num_class, num_boxes = scores.shape
 
     scores = reshape(scores, (batch * num_class, num_boxes))
     sorted_scores, sorted_indices = _dispatch_sort(scores, ret_type="both")
     valid_count = _get_valid_box_count(sorted_scores, score_threshold)
 
-    selected_indices, num_detections = run_all_class_nms(
+    selected_indices, selected_scores, num_detections = run_all_class_nms(
         boxes,
         sorted_scores,
         sorted_indices,
         valid_count,
         max_output_boxes_per_class,
         iou_threshold,
         _nms_loop,
+        return_scores=(output_format == "tensorflow"),
     )
 
+    if output_format == "onnx":
+        row_offsets, num_total_detections = exclusive_scan(
+            num_detections, return_reduction=True, output_dtype="int64"
+        )
+        selected_indices = collect_selected_indices(
+            num_class, selected_indices, num_detections, row_offsets, _collect_selected_indices_ir
+        )
+        return [selected_indices, num_total_detections]
+
+    num_detections_per_batch = reshape(num_detections, (batch, num_class))
     row_offsets, num_total_detections = exclusive_scan(
-        num_detections, return_reduction=True, output_dtype="int64"
+        num_detections_per_batch, return_reduction=True, output_dtype="int64", axis=1

Review comment:
       I see. I didn't hit this error because I was not using thrust, and this code path is only for thrust. Fixed by adding `cast(0, reduction.dtype)`.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-855097348


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


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi edited a comment on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854239897


   lol thrust is the bottleneck now? Note that our sorting performance improved thanks to https://github.com/apache/tvm/pull/7611, 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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644302217



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,92 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    tile_batch_reps = (
+        _op.concatenate([batch_size, 1], axis=0)

Review comment:
       oops you are right. Actually the existing converter does not seem to support both dynamic batch and num boxes either: https://github.com/apache/tvm/blob/6baccc1305b13f18ce8ac5db5fcbbadb8c278454/python/tvm/relay/frontend/tensorflow.py#L827-L828
   
   So I'll just add an assertion to check if batch_size and num_boxes are integers and raise an error otherwise.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-855112421


   > @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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644337166



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,89 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    assert isinstance(batch_size, int), "dynamic batch size not supported yet."
+    tile_batch_reps = _op.const([batch_size, 1])
+    box_range_2d = _op.tile(box_range, tile_batch_reps)
+    valid_mask = _op.cast(
+        _op.less(box_range_2d, _op.expand_dims(num_detections, axis=1)), "float32"
+    )
+
+    def select_topk(do_zero_pad):
+        def true_branch():
+            arange = _op.arange(
+                _op.const(0, dtype="int64"),
+                _op.const(max_output_boxes_per_batch, dtype="int64"),
+                dtype="int64",
+            )
+            pad = _op.full(
+                _op.const(0, dtype="int64"), (max_total_size - max_output_boxes_per_batch,)
+            )
+            topk_indices = _op.tile(_op.concatenate([arange, pad], 0), tile_batch_reps)
+            nmsed_scores = _op.gather(selected_scores, 1, topk_indices)
+            nmsed_scores = nmsed_scores * valid_mask
+            return nmsed_scores, topk_indices
+
+        def false_branch():
+            if isinstance(max_output_boxes_per_class, int):
+                # Do topk on smaller input if possible
+                # TODO(masahi): use axes argument in strided slice when it becomes available
+                slice_mx = _op.const([-1, max_output_boxes_per_class * num_class], dtype="int64")
+                selected_scores_slice = _op.strided_slice(
+                    selected_scores, begin=_op.const([0, 0], dtype="int64"), end=slice_mx

Review comment:
       should be fixed




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644342291



##########
File path: python/tvm/topi/cuda/nms.py
##########
@@ -1012,39 +1079,69 @@ def all_class_non_max_suppression(
     score_threshold : float or tvm.te.Tensor, optional
         Score threshold to filter out low score boxes early
 
+    output_format : str, optional
+        "onnx" or "tensorflow", see below
+
     Returns
     -------
-    out : [tvm.te.Tensor, tvm.te.Tensor]
-        The output is two tensors, the first is `indices` of size
+    out : list of tvm.te.Tensor
+        If `output_format` is "onnx", the output is two tensors. The first is `indices` of size
         `(batch_size * num_class* num_boxes , 3)` and the second is a scalar tensor
         `num_total_detection` of shape `(1,)` representing the total number of selected
-        boxes. Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
+        boxes. The three values in `indices` encode batch, class, and box indices.
+        Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
         first, in descending of scores, followed by boxes from batch 0, class 1 etc. Out of
         `batch_size * num_class* num_boxes` rows of indices, only the first `num_total_detection`
         rows are valid.
+
+        If `output_format` is "tensorflow", the output is three tensors, the first
+        is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the second is `scores` of
+        size `(batch_size, num_class * num_boxes)`, and the third is `num_total_detection` of size
+        `(batch_size,)` representing the total number of selected boxes per batch. The two values
+        in `indices` encode class and box indices. Of num_class * num_boxes boxes in `indices` at
+        batch b, only the first `num_total_detection[b]` entries are valid. The second axis of
+        `indices` and `scores` are sorted within each class by box scores, but not across classes.
+        So the box indices and scores for the class 0 come first in a sorted order, followed by
+        the class 1 etc.
     """
     batch, num_class, num_boxes = scores.shape
 
     scores = reshape(scores, (batch * num_class, num_boxes))
     sorted_scores, sorted_indices = _dispatch_sort(scores, ret_type="both")
     valid_count = _get_valid_box_count(sorted_scores, score_threshold)
 
-    selected_indices, num_detections = run_all_class_nms(
+    selected_indices, selected_scores, num_detections = run_all_class_nms(
         boxes,
         sorted_scores,
         sorted_indices,
         valid_count,
         max_output_boxes_per_class,
         iou_threshold,
         _nms_loop,
+        return_scores=(output_format == "tensorflow"),
     )
 
+    if output_format == "onnx":
+        row_offsets, num_total_detections = exclusive_scan(
+            num_detections, return_reduction=True, output_dtype="int64"
+        )
+        selected_indices = collect_selected_indices(
+            num_class, selected_indices, num_detections, row_offsets, _collect_selected_indices_ir
+        )
+        return [selected_indices, num_total_detections]
+
+    num_detections_per_batch = reshape(num_detections, (batch, num_class))
     row_offsets, num_total_detections = exclusive_scan(
-        num_detections, return_reduction=True, output_dtype="int64"
+        num_detections_per_batch, return_reduction=True, output_dtype="int64", axis=1

Review comment:
       Nice, thanks!
   




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644337398



##########
File path: python/tvm/topi/cuda/nms.py
##########
@@ -1012,39 +1079,69 @@ def all_class_non_max_suppression(
     score_threshold : float or tvm.te.Tensor, optional
         Score threshold to filter out low score boxes early
 
+    output_format : str, optional
+        "onnx" or "tensorflow", see below
+
     Returns
     -------
-    out : [tvm.te.Tensor, tvm.te.Tensor]
-        The output is two tensors, the first is `indices` of size
+    out : list of tvm.te.Tensor
+        If `output_format` is "onnx", the output is two tensors. The first is `indices` of size
         `(batch_size * num_class* num_boxes , 3)` and the second is a scalar tensor
         `num_total_detection` of shape `(1,)` representing the total number of selected
-        boxes. Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
+        boxes. The three values in `indices` encode batch, class, and box indices.
+        Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
         first, in descending of scores, followed by boxes from batch 0, class 1 etc. Out of
         `batch_size * num_class* num_boxes` rows of indices, only the first `num_total_detection`
         rows are valid.
+
+        If `output_format` is "tensorflow", the output is three tensors, the first
+        is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the second is `scores` of
+        size `(batch_size, num_class * num_boxes)`, and the third is `num_total_detection` of size
+        `(batch_size,)` representing the total number of selected boxes per batch. The two values
+        in `indices` encode class and box indices. Of num_class * num_boxes boxes in `indices` at
+        batch b, only the first `num_total_detection[b]` entries are valid. The second axis of
+        `indices` and `scores` are sorted within each class by box scores, but not across classes.
+        So the box indices and scores for the class 0 come first in a sorted order, followed by
+        the class 1 etc.
     """
     batch, num_class, num_boxes = scores.shape
 
     scores = reshape(scores, (batch * num_class, num_boxes))
     sorted_scores, sorted_indices = _dispatch_sort(scores, ret_type="both")
     valid_count = _get_valid_box_count(sorted_scores, score_threshold)
 
-    selected_indices, num_detections = run_all_class_nms(
+    selected_indices, selected_scores, num_detections = run_all_class_nms(
         boxes,
         sorted_scores,
         sorted_indices,
         valid_count,
         max_output_boxes_per_class,
         iou_threshold,
         _nms_loop,
+        return_scores=(output_format == "tensorflow"),
     )
 
+    if output_format == "onnx":
+        row_offsets, num_total_detections = exclusive_scan(
+            num_detections, return_reduction=True, output_dtype="int64"
+        )
+        selected_indices = collect_selected_indices(
+            num_class, selected_indices, num_detections, row_offsets, _collect_selected_indices_ir
+        )
+        return [selected_indices, num_total_detections]
+
+    num_detections_per_batch = reshape(num_detections, (batch, num_class))
     row_offsets, num_total_detections = exclusive_scan(
-        num_detections, return_reduction=True, output_dtype="int64"
+        num_detections_per_batch, return_reduction=True, output_dtype="int64", axis=1

Review comment:
       I see. I didn't hit this error because I was not using thrust, and this code path is only for thrust. Fixed by adding `cast(0, reduction.dtype)`.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi edited a comment on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-853506007


   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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644331266



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,89 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    assert isinstance(batch_size, int), "dynamic batch size not supported yet."
+    tile_batch_reps = _op.const([batch_size, 1])
+    box_range_2d = _op.tile(box_range, tile_batch_reps)
+    valid_mask = _op.cast(
+        _op.less(box_range_2d, _op.expand_dims(num_detections, axis=1)), "float32"
+    )
+
+    def select_topk(do_zero_pad):
+        def true_branch():
+            arange = _op.arange(
+                _op.const(0, dtype="int64"),
+                _op.const(max_output_boxes_per_batch, dtype="int64"),
+                dtype="int64",
+            )
+            pad = _op.full(
+                _op.const(0, dtype="int64"), (max_total_size - max_output_boxes_per_batch,)
+            )
+            topk_indices = _op.tile(_op.concatenate([arange, pad], 0), tile_batch_reps)
+            nmsed_scores = _op.gather(selected_scores, 1, topk_indices)
+            nmsed_scores = nmsed_scores * valid_mask
+            return nmsed_scores, topk_indices
+
+        def false_branch():
+            if isinstance(max_output_boxes_per_class, int):
+                # Do topk on smaller input if possible
+                # TODO(masahi): use axes argument in strided slice when it becomes available
+                slice_mx = _op.const([-1, max_output_boxes_per_class * num_class], dtype="int64")
+                selected_scores_slice = _op.strided_slice(
+                    selected_scores, begin=_op.const([0, 0], dtype="int64"), end=slice_mx

Review comment:
       I see, I'd rather use `axes` argument introduced in https://github.com/apache/tvm/pull/8165 after it is merged. This saves us from specifying the begin and end for the batch dim. I was using this in my dev branch, and realized that `axes` argument is not in `main` yet, so added this fishy workaround. And tests in `test_forward.py` have `max_output_boxes_per_class` as a variable, so we don't hit this code path.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854239897


   lol thrust is the bottleneck now? 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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644321057



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,89 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    assert isinstance(batch_size, int), "dynamic batch size not supported yet."
+    tile_batch_reps = _op.const([batch_size, 1])
+    box_range_2d = _op.tile(box_range, tile_batch_reps)
+    valid_mask = _op.cast(
+        _op.less(box_range_2d, _op.expand_dims(num_detections, axis=1)), "float32"
+    )
+
+    def select_topk(do_zero_pad):
+        def true_branch():
+            arange = _op.arange(
+                _op.const(0, dtype="int64"),
+                _op.const(max_output_boxes_per_batch, dtype="int64"),
+                dtype="int64",
+            )
+            pad = _op.full(
+                _op.const(0, dtype="int64"), (max_total_size - max_output_boxes_per_batch,)
+            )
+            topk_indices = _op.tile(_op.concatenate([arange, pad], 0), tile_batch_reps)
+            nmsed_scores = _op.gather(selected_scores, 1, topk_indices)
+            nmsed_scores = nmsed_scores * valid_mask
+            return nmsed_scores, topk_indices
+
+        def false_branch():
+            if isinstance(max_output_boxes_per_class, int):
+                # Do topk on smaller input if possible
+                # TODO(masahi): use axes argument in strided slice when it becomes available
+                slice_mx = _op.const([-1, max_output_boxes_per_class * num_class], dtype="int64")
+                selected_scores_slice = _op.strided_slice(
+                    selected_scores, begin=_op.const([0, 0], dtype="int64"), end=slice_mx

Review comment:
       Need to use `slice_mode="size"` for this strided_slice, otherwise using -1 for end is returning a tensor with batch size of 0.
   
   ```
   %988 = vision.all_class_non_max_suppression(%986, %987, 100 /* ty=int32 */, 0.6f /* ty=float32 */, 0.3f /* ty=float32 */, meta[relay.attrs.AllClassNonMaximumSuppressionAttrs][0]) /* ty=(Tensor[(1, 172530, 2), int64], Tensor[(1, 172530), float32], Tensor[(1), int64]) */;
   %989 = %988.1;
   %990 = strided_slice(%989, begin=[0, 0], end=[-1, 9000], strides=[1]) /* ty=Tensor[(0, 9000), float32] */;
   ```

##########
File path: python/tvm/topi/cuda/nms.py
##########
@@ -1012,39 +1079,69 @@ def all_class_non_max_suppression(
     score_threshold : float or tvm.te.Tensor, optional
         Score threshold to filter out low score boxes early
 
+    output_format : str, optional
+        "onnx" or "tensorflow", see below
+
     Returns
     -------
-    out : [tvm.te.Tensor, tvm.te.Tensor]
-        The output is two tensors, the first is `indices` of size
+    out : list of tvm.te.Tensor
+        If `output_format` is "onnx", the output is two tensors. The first is `indices` of size
         `(batch_size * num_class* num_boxes , 3)` and the second is a scalar tensor
         `num_total_detection` of shape `(1,)` representing the total number of selected
-        boxes. Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
+        boxes. The three values in `indices` encode batch, class, and box indices.
+        Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
         first, in descending of scores, followed by boxes from batch 0, class 1 etc. Out of
         `batch_size * num_class* num_boxes` rows of indices, only the first `num_total_detection`
         rows are valid.
+
+        If `output_format` is "tensorflow", the output is three tensors, the first
+        is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the second is `scores` of
+        size `(batch_size, num_class * num_boxes)`, and the third is `num_total_detection` of size
+        `(batch_size,)` representing the total number of selected boxes per batch. The two values
+        in `indices` encode class and box indices. Of num_class * num_boxes boxes in `indices` at
+        batch b, only the first `num_total_detection[b]` entries are valid. The second axis of
+        `indices` and `scores` are sorted within each class by box scores, but not across classes.
+        So the box indices and scores for the class 0 come first in a sorted order, followed by
+        the class 1 etc.
     """
     batch, num_class, num_boxes = scores.shape
 
     scores = reshape(scores, (batch * num_class, num_boxes))
     sorted_scores, sorted_indices = _dispatch_sort(scores, ret_type="both")
     valid_count = _get_valid_box_count(sorted_scores, score_threshold)
 
-    selected_indices, num_detections = run_all_class_nms(
+    selected_indices, selected_scores, num_detections = run_all_class_nms(
         boxes,
         sorted_scores,
         sorted_indices,
         valid_count,
         max_output_boxes_per_class,
         iou_threshold,
         _nms_loop,
+        return_scores=(output_format == "tensorflow"),
     )
 
+    if output_format == "onnx":
+        row_offsets, num_total_detections = exclusive_scan(
+            num_detections, return_reduction=True, output_dtype="int64"
+        )
+        selected_indices = collect_selected_indices(
+            num_class, selected_indices, num_detections, row_offsets, _collect_selected_indices_ir
+        )
+        return [selected_indices, num_total_detections]
+
+    num_detections_per_batch = reshape(num_detections, (batch, num_class))
     row_offsets, num_total_detections = exclusive_scan(
-        num_detections, return_reduction=True, output_dtype="int64"
+        num_detections_per_batch, return_reduction=True, output_dtype="int64", axis=1

Review comment:
       Because of the int64 type, I had to set this line
   https://github.com/apache/tvm/blob/main/python/tvm/topi/cuda/scan.py#L234
   
   to `reduction[tid] = cast(0, "int64")`
   
   Otherwise I got the error: `data type does not match content type int32 vs int64`




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m edited a comment on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m edited a comment on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854235480


   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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854235480


   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 ms |
   
   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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m commented on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m commented on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854235480


   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 ms |
   
   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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m edited a comment on pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m edited a comment on pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#issuecomment-854235480


   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.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] masahi commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644306389



##########
File path: python/tvm/topi/transform.py
##########
@@ -483,7 +483,7 @@ def gather(data, axis, indices):
     return cpp.gather(data, axis, indices)
 
 
-def gather_nd(a, indices):
+def gather_nd(a, indices, batch_dims=0):

Review comment:
       I was using this before but not anymore. Reverted.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m commented on a change in pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m commented on a change in pull request #8174:
URL: https://github.com/apache/tvm/pull/8174#discussion_r644321057



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,89 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    assert isinstance(batch_size, int), "dynamic batch size not supported yet."
+    tile_batch_reps = _op.const([batch_size, 1])
+    box_range_2d = _op.tile(box_range, tile_batch_reps)
+    valid_mask = _op.cast(
+        _op.less(box_range_2d, _op.expand_dims(num_detections, axis=1)), "float32"
+    )
+
+    def select_topk(do_zero_pad):
+        def true_branch():
+            arange = _op.arange(
+                _op.const(0, dtype="int64"),
+                _op.const(max_output_boxes_per_batch, dtype="int64"),
+                dtype="int64",
+            )
+            pad = _op.full(
+                _op.const(0, dtype="int64"), (max_total_size - max_output_boxes_per_batch,)
+            )
+            topk_indices = _op.tile(_op.concatenate([arange, pad], 0), tile_batch_reps)
+            nmsed_scores = _op.gather(selected_scores, 1, topk_indices)
+            nmsed_scores = nmsed_scores * valid_mask
+            return nmsed_scores, topk_indices
+
+        def false_branch():
+            if isinstance(max_output_boxes_per_class, int):
+                # Do topk on smaller input if possible
+                # TODO(masahi): use axes argument in strided slice when it becomes available
+                slice_mx = _op.const([-1, max_output_boxes_per_class * num_class], dtype="int64")
+                selected_scores_slice = _op.strided_slice(
+                    selected_scores, begin=_op.const([0, 0], dtype="int64"), end=slice_mx

Review comment:
       Need to use `slice_mode="size"` for this strided_slice, otherwise using -1 for end is returning a tensor with batch size of 0.
   
   ```
   %988 = vision.all_class_non_max_suppression(%986, %987, 100 /* ty=int32 */, 0.6f /* ty=float32 */, 0.3f /* ty=float32 */, meta[relay.attrs.AllClassNonMaximumSuppressionAttrs][0]) /* ty=(Tensor[(1, 172530, 2), int64], Tensor[(1, 172530), float32], Tensor[(1), int64]) */;
   %989 = %988.1;
   %990 = strided_slice(%989, begin=[0, 0], end=[-1, 9000], strides=[1]) /* ty=Tensor[(0, 9000), float32] */;
   ```

##########
File path: python/tvm/topi/cuda/nms.py
##########
@@ -1012,39 +1079,69 @@ def all_class_non_max_suppression(
     score_threshold : float or tvm.te.Tensor, optional
         Score threshold to filter out low score boxes early
 
+    output_format : str, optional
+        "onnx" or "tensorflow", see below
+
     Returns
     -------
-    out : [tvm.te.Tensor, tvm.te.Tensor]
-        The output is two tensors, the first is `indices` of size
+    out : list of tvm.te.Tensor
+        If `output_format` is "onnx", the output is two tensors. The first is `indices` of size
         `(batch_size * num_class* num_boxes , 3)` and the second is a scalar tensor
         `num_total_detection` of shape `(1,)` representing the total number of selected
-        boxes. Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
+        boxes. The three values in `indices` encode batch, class, and box indices.
+        Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
         first, in descending of scores, followed by boxes from batch 0, class 1 etc. Out of
         `batch_size * num_class* num_boxes` rows of indices, only the first `num_total_detection`
         rows are valid.
+
+        If `output_format` is "tensorflow", the output is three tensors, the first
+        is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the second is `scores` of
+        size `(batch_size, num_class * num_boxes)`, and the third is `num_total_detection` of size
+        `(batch_size,)` representing the total number of selected boxes per batch. The two values
+        in `indices` encode class and box indices. Of num_class * num_boxes boxes in `indices` at
+        batch b, only the first `num_total_detection[b]` entries are valid. The second axis of
+        `indices` and `scores` are sorted within each class by box scores, but not across classes.
+        So the box indices and scores for the class 0 come first in a sorted order, followed by
+        the class 1 etc.
     """
     batch, num_class, num_boxes = scores.shape
 
     scores = reshape(scores, (batch * num_class, num_boxes))
     sorted_scores, sorted_indices = _dispatch_sort(scores, ret_type="both")
     valid_count = _get_valid_box_count(sorted_scores, score_threshold)
 
-    selected_indices, num_detections = run_all_class_nms(
+    selected_indices, selected_scores, num_detections = run_all_class_nms(
         boxes,
         sorted_scores,
         sorted_indices,
         valid_count,
         max_output_boxes_per_class,
         iou_threshold,
         _nms_loop,
+        return_scores=(output_format == "tensorflow"),
     )
 
+    if output_format == "onnx":
+        row_offsets, num_total_detections = exclusive_scan(
+            num_detections, return_reduction=True, output_dtype="int64"
+        )
+        selected_indices = collect_selected_indices(
+            num_class, selected_indices, num_detections, row_offsets, _collect_selected_indices_ir
+        )
+        return [selected_indices, num_total_detections]
+
+    num_detections_per_batch = reshape(num_detections, (batch, num_class))
     row_offsets, num_total_detections = exclusive_scan(
-        num_detections, return_reduction=True, output_dtype="int64"
+        num_detections_per_batch, return_reduction=True, output_dtype="int64", axis=1

Review comment:
       Because of the int64 type, I had to set this line
   https://github.com/apache/tvm/blob/main/python/tvm/topi/cuda/scan.py#L234
   
   to `reduction[tid] = cast(0, "int64")`
   
   Otherwise I got the error: `data type does not match content type int32 vs int64`

##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -793,6 +793,89 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def convert_combined_nms_with_all_class_nms(
+    batch_size,
+    max_output_boxes_per_batch,
+    num_class,
+    boxes,
+    scores,
+    max_output_boxes_per_class,
+    iou_threshold,
+    score_threshold,
+    max_total_size,
+    clip_boxes,
+):
+    """Converts TF combined_nms using Relay all_class_max_suppression op"""
+    (selected_indices, selected_scores, num_detections,) = _op.vision.all_class_non_max_suppression(
+        boxes,
+        scores,
+        max_output_boxes_per_class,
+        iou_threshold,
+        score_threshold,
+        output_format="tensorflow",
+    )
+    box_range = _op.arange(
+        _op.const(0, dtype="int64"), _op.const(max_total_size, dtype="int64"), dtype="int64"
+    )
+    assert isinstance(batch_size, int), "dynamic batch size not supported yet."
+    tile_batch_reps = _op.const([batch_size, 1])
+    box_range_2d = _op.tile(box_range, tile_batch_reps)
+    valid_mask = _op.cast(
+        _op.less(box_range_2d, _op.expand_dims(num_detections, axis=1)), "float32"
+    )
+
+    def select_topk(do_zero_pad):
+        def true_branch():
+            arange = _op.arange(
+                _op.const(0, dtype="int64"),
+                _op.const(max_output_boxes_per_batch, dtype="int64"),
+                dtype="int64",
+            )
+            pad = _op.full(
+                _op.const(0, dtype="int64"), (max_total_size - max_output_boxes_per_batch,)
+            )
+            topk_indices = _op.tile(_op.concatenate([arange, pad], 0), tile_batch_reps)
+            nmsed_scores = _op.gather(selected_scores, 1, topk_indices)
+            nmsed_scores = nmsed_scores * valid_mask
+            return nmsed_scores, topk_indices
+
+        def false_branch():
+            if isinstance(max_output_boxes_per_class, int):
+                # Do topk on smaller input if possible
+                # TODO(masahi): use axes argument in strided slice when it becomes available
+                slice_mx = _op.const([-1, max_output_boxes_per_class * num_class], dtype="int64")
+                selected_scores_slice = _op.strided_slice(
+                    selected_scores, begin=_op.const([0, 0], dtype="int64"), end=slice_mx

Review comment:
       Thanks for fixing!

##########
File path: python/tvm/topi/cuda/nms.py
##########
@@ -1012,39 +1079,69 @@ def all_class_non_max_suppression(
     score_threshold : float or tvm.te.Tensor, optional
         Score threshold to filter out low score boxes early
 
+    output_format : str, optional
+        "onnx" or "tensorflow", see below
+
     Returns
     -------
-    out : [tvm.te.Tensor, tvm.te.Tensor]
-        The output is two tensors, the first is `indices` of size
+    out : list of tvm.te.Tensor
+        If `output_format` is "onnx", the output is two tensors. The first is `indices` of size
         `(batch_size * num_class* num_boxes , 3)` and the second is a scalar tensor
         `num_total_detection` of shape `(1,)` representing the total number of selected
-        boxes. Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
+        boxes. The three values in `indices` encode batch, class, and box indices.
+        Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
         first, in descending of scores, followed by boxes from batch 0, class 1 etc. Out of
         `batch_size * num_class* num_boxes` rows of indices, only the first `num_total_detection`
         rows are valid.
+
+        If `output_format` is "tensorflow", the output is three tensors, the first
+        is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the second is `scores` of
+        size `(batch_size, num_class * num_boxes)`, and the third is `num_total_detection` of size
+        `(batch_size,)` representing the total number of selected boxes per batch. The two values
+        in `indices` encode class and box indices. Of num_class * num_boxes boxes in `indices` at
+        batch b, only the first `num_total_detection[b]` entries are valid. The second axis of
+        `indices` and `scores` are sorted within each class by box scores, but not across classes.
+        So the box indices and scores for the class 0 come first in a sorted order, followed by
+        the class 1 etc.
     """
     batch, num_class, num_boxes = scores.shape
 
     scores = reshape(scores, (batch * num_class, num_boxes))
     sorted_scores, sorted_indices = _dispatch_sort(scores, ret_type="both")
     valid_count = _get_valid_box_count(sorted_scores, score_threshold)
 
-    selected_indices, num_detections = run_all_class_nms(
+    selected_indices, selected_scores, num_detections = run_all_class_nms(
         boxes,
         sorted_scores,
         sorted_indices,
         valid_count,
         max_output_boxes_per_class,
         iou_threshold,
         _nms_loop,
+        return_scores=(output_format == "tensorflow"),
     )
 
+    if output_format == "onnx":
+        row_offsets, num_total_detections = exclusive_scan(
+            num_detections, return_reduction=True, output_dtype="int64"
+        )
+        selected_indices = collect_selected_indices(
+            num_class, selected_indices, num_detections, row_offsets, _collect_selected_indices_ir
+        )
+        return [selected_indices, num_total_detections]
+
+    num_detections_per_batch = reshape(num_detections, (batch, num_class))
     row_offsets, num_total_detections = exclusive_scan(
-        num_detections, return_reduction=True, output_dtype="int64"
+        num_detections_per_batch, return_reduction=True, output_dtype="int64", axis=1

Review comment:
       Nice, thanks!
   




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] trevor-m merged pull request #8174: [Relay, TF] Support converting TF combined_nms using Relay all_class_nms

Posted by GitBox <gi...@apache.org>.
trevor-m merged pull request #8174:
URL: https://github.com/apache/tvm/pull/8174


   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org