minor refactor

masahi · May 29, 2021 · f88e2a3 · f88e2a3
1 parent f2d7ed4
commit f88e2a3
Showing 1 changed file with 8 additions and 10 deletions.
diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
@@ -796,16 +796,14 @@ def _impl(inputs, attr, params, mod):
 def _combined_nms():
     def all_class_impl(
         batch_size,
-        num_boxes,
-        num_classes,
+        max_output_boxes_per_batch,
         boxes,
         scores,
         max_output_boxes_per_class,
         iou_threshold,
         score_threshold,
         max_total_size,
         clip_boxes,
-        mod,
     ):
         (
             selected_indices,
@@ -835,14 +833,15 @@ def all_class_impl(
         )
 
         # TODO: support dynamic num_boxes
-        max_output_boxes = num_boxes * num_classes
-        if max_output_boxes < max_total_size:
+        if max_output_boxes_per_batch < max_total_size:
             arange = _op.arange(
                 _op.const(0, dtype="int64"),
-                _op.const(max_output_boxes, 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,))
+            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
@@ -898,18 +897,17 @@ def _impl(inputs, attr, params, mod):
         if q == 1:
             boxes = _op.squeeze(boxes, axis=[2])
             scores_trans = _op.transpose(scores, [0, 2, 1])
+            max_output_boxes_per_batch = num_anchors * num_classes
             return all_class_impl(
                 batch_size,
-                num_anchors,
-                num_classes,
+                max_output_boxes_per_batch,
                 boxes,
                 scores_trans,
                 max_output_size,
                 iou_threshold,
                 score_threshold,
                 max_total_size.data.numpy().item(),
                 attr["clip_boxes"],
-                mod,
             )
 
         boxes = _op.reshape(boxes, newshape=[batch_size, num_anchors * num_classes, 4])