use topk instead of sort

Summary: now pytorch/pytorch#22812 is fixed Reviewed By: zhanghang1989 Differential Revision: D32610251 fbshipit-source-id: e099a2c53f71cca95af35aafc26ab59f9613c07b
yusiyoh · Nov 27, 2021 · 1ad5759 · 1ad5759
1 parent 4606450
commit 1ad5759
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 14 deletions.
diff --git a/detectron2/modeling/meta_arch/dense_detector.py b/detectron2/modeling/meta_arch/dense_detector.py
@@ -208,10 +208,8 @@ def _decode_per_level_predictions(
 
         # 2. Keep top k top scoring boxes only
         num_topk = min(topk_candidates, topk_idxs.size(0))
-        # torch.sort is actually faster than .topk (https://github.com/pytorch/pytorch/issues/22812)
-        pred_scores, idxs = pred_scores.sort(descending=True)
-        pred_scores = pred_scores[:num_topk]
-        topk_idxs = topk_idxs[idxs[:num_topk]]
+        pred_scores, idxs = pred_scores.topk(num_topk)
+        topk_idxs = topk_idxs[idxs]
 
         anchor_idxs, classes_idxs = topk_idxs.unbind(dim=1)
 

diff --git a/detectron2/modeling/proposal_generator/proposal_utils.py b/detectron2/modeling/proposal_generator/proposal_utils.py
@@ -72,11 +72,7 @@ def find_top_rpn_proposals(
         else:
             num_proposals_i = min(Hi_Wi_A, pre_nms_topk)
 
-        # sort is faster than topk: https://github.com/pytorch/pytorch/issues/22812
-        # topk_scores_i, topk_idx = logits_i.topk(num_proposals_i, dim=1)
-        logits_i, idx = logits_i.sort(descending=True, dim=1)
-        topk_scores_i = logits_i.narrow(1, 0, num_proposals_i)
-        topk_idx = idx.narrow(1, 0, num_proposals_i)
+        topk_scores_i, topk_idx = logits_i.topk(num_proposals_i, dim=1)
 
         # each is N x topk
         topk_proposals_i = proposals_i[batch_idx[:, None], topk_idx]  # N x topk x 4

diff --git a/detectron2/modeling/proposal_generator/rrpn.py b/detectron2/modeling/proposal_generator/rrpn.py
@@ -73,11 +73,7 @@ def find_top_rrpn_proposals(
         else:
             num_proposals_i = min(Hi_Wi_A, pre_nms_topk)
 
-        # sort is faster than topk (https://github.com/pytorch/pytorch/issues/22812)
-        # topk_scores_i, topk_idx = logits_i.topk(num_proposals_i, dim=1)
-        logits_i, idx = logits_i.sort(descending=True, dim=1)
-        topk_scores_i = logits_i[batch_idx, :num_proposals_i]
-        topk_idx = idx[batch_idx, :num_proposals_i]
+        topk_scores_i, topk_idx = logits_i.topk(num_proposals_i, dim=1)
 
         # each is N x topk
         topk_proposals_i = proposals_i[batch_idx[:, None], topk_idx]  # N x topk x 5