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:
[email protected]
