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

[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