Supporting non maximum suppresion op with dynamic output shape

masahi · February 27, 2020, 7:28am

Hi, I’m trying to translate PyTorch’s NMS op below, whose output shape depends on input:

pytorch/vision/blob/2e8bcf8ba48ea94e30a362e1a57e5e3de8dbd2d0/torchvision/ops/boxes.py#L29-L36


Returns
-------
keep : Tensor
    int64 tensor with the indices
    of the elements that have been kept
    by NMS, sorted in decreasing order of scores
"""
return torch.ops.torchvision.nms(boxes, scores, iou_threshold)

Can we translate this to Relay using its recent dynamic feature? I’m not familiar with shape function etc.

@haichen @zhiics @jroesch

haichen · February 28, 2020, 6:04am

Currently it’s not very easy to add shape function for the nms op because computing the output shape basically has the same cost of computing nms. The better way to solve this is to add an extra mode, upper bound, to the shape function that only calculates the upper bound of output shape. And then resize it after nms computation finishes.