kevinthesun commented on a change in pull request #4312: URL: https://github.com/apache/incubator-tvm/pull/4312#discussion_r411062395
########## File path: topi/python/topi/vision/nms.py ########## @@ -64,7 +67,56 @@ def hybrid_rearrange_out(data, one): @hybrid.script -def hybrid_get_valid_counts(data, score_threshold, id_index, score_index, one): +def hybrid_rearrange_indices_out(data, one, batch_size): + """Hybrid routine to rearrange nms output to + move all valid entries to top. + + Parameters + ---------- + data : tvm.te.Tensor or numpy NDArray + NMS output. 3-D tensor with shape + [batch_size, num_anchors, 6] or + [batch_size, num_anchors, 5], or 2-D + tensor with shape [batch_size, num_anchors]. + + one: tvm.tir.const + Constant one with the same dtype as data. + + batch_size: tvm.tir.IntImm or tvm.tir.Var + Batch size. We need to pass it in since hybrid script doesn't support + binding variable to symbolic dim. + + Returns + ------- + output : tvm.te.Tensor or numpy NDArray + 2-D tensor with shape [batch_size, num_anchors]. + + valid_box_count : tvm.te.Tensor or numpy NDArray + Tensor with shape [batch_size, 1], indicates + the valid number of boxes. + """ + num_anchors = data.shape[1] + valid_box_count = output_tensor((batch_size, 1), "int32") + output = output_tensor((batch_size, num_anchors), data.dtype) + + for i in parallel(batch_size): + valid_idx = 0 + for j in range(num_anchors): + if data[i, j] >= 0: + output[i, valid_idx] = data[i, j] + valid_idx += 1 + if data[i, j] > num_anchors or data[i, j] < -num_anchors: + output[i, valid_idx] = 0 + valid_idx += 1 + if j >= valid_idx: + output[i, j] = -one + valid_box_count[i, 0] = valid_idx + + return output, valid_box_count + + [email protected] +def hybrid_get_valid_counts(data, score_threshold, id_index, score_index, one, batch_size): Review comment: Same. ---------------------------------------------------------------- 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]
