call batched_nms from torchvision directly

Summary:
torchvision already has the thresholding logic.
The threshold in torchvision is not good, which will be fixed in pytorch/vision#4990

Reviewed By: zhanghang1989

Differential Revision: D32661585

fbshipit-source-id: 971b68a914c2bba29dff37298c38e31b02735c1d

Author: ppwwyyxx

detectron2/layers/nms.py

@@ -1,33 +1,23 @@
 # -*- coding: utf-8 -*-
 # Copyright (c) Facebook, Inc. and its affiliates.
 
-from typing import List
 import torch
 from torchvision.ops import boxes as box_ops
-from torchvision.ops import nms  # BC-compat
+from torchvision.ops import nms  # noqa . for compatibility
 
 
 def batched_nms(
     boxes: torch.Tensor, scores: torch.Tensor, idxs: torch.Tensor, iou_threshold: float
 ):
     """
-    Same as torchvision.ops.boxes.batched_nms, but safer.
+    Same as torchvision.ops.boxes.batched_nms, but with float().
     """
     assert boxes.shape[-1] == 4
-    # TODO may need better strategy.
-    # Investigate after having a fully-cuda NMS op.
-    if len(boxes) < 40000:
-        # fp16 does not have enough range for batched NMS
-        return box_ops.batched_nms(boxes.float(), scores, idxs, iou_threshold)
-
-    result_mask = scores.new_zeros(scores.size(), dtype=torch.bool)
-    for id in torch.jit.annotate(List[int], torch.unique(idxs).cpu().tolist()):
-        mask = (idxs == id).nonzero().view(-1)
-        keep = nms(boxes[mask], scores[mask], iou_threshold)
-        result_mask[mask[keep]] = True
-    keep = result_mask.nonzero().view(-1)
-    keep = keep[scores[keep].argsort(descending=True)]
-    return keep
+    # Note: Torchvision already has a strategy (https://github.com/pytorch/vision/issues/1311)
+    # to decide whether to use coordinate trick or for loop to implement batched_nms. So we
+    # just call it directly.
+    # Fp16 does not have enough range for batched NMS, so adding float().
+    return box_ops.batched_nms(boxes.float(), scores, idxs, iou_threshold)
 
 
 # Note: this function (nms_rotated) might be moved into