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