From 2b76db991717c23c10f0eb51d443fc7da245a44d Mon Sep 17 00:00:00 2001
From: RichardWooSJTU <37864677+RichardWooSJTU@users.noreply.github.com>
Date: Mon, 19 Sep 2022 10:52:24 +0800
Subject: [PATCH] [vision.ops.nms] Fix return order error and duplicate results
 with specific inputs (#46148)

* fix return order error and duplicate results with specific inputs
---
 paddle/phi/infermeta/unary.cc                 |  4 ++--
 paddle/phi/kernels/cpu/nms_kernel.cc          | 24 ++++++++++++++-----
 paddle/phi/kernels/gpu/nms_kernel.cu          |  5 ++--
 .../fluid/tests/unittests/test_nms_op.py      |  2 +-
 python/paddle/vision/ops.py                   |  4 +++-
 5 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
index 39db2579ecb..2fdb32644ad 100644
--- a/paddle/phi/infermeta/unary.cc
+++ b/paddle/phi/infermeta/unary.cc
@@ -2035,8 +2035,8 @@ void NMSInferMeta(const MetaTensor& x, float threshold, MetaTensor* out) {
                         "whose shape must be [N, 4] "
                         "N is the number of boxes "
                         "in last dimension in format [x1, x2, y1, y2]. "));
-  auto num_boxes = boxes_dim[0];
-  out->set_dims(phi::make_ddim({num_boxes}));
+  out->set_dims(phi::make_ddim({-1}));
+  out->set_dtype(DataType::INT64);
 }
 
 void NormInferMeta(const MetaTensor& x,
diff --git a/paddle/phi/kernels/cpu/nms_kernel.cc b/paddle/phi/kernels/cpu/nms_kernel.cc
index 7e656b14f1f..4b56f6bb951 100644
--- a/paddle/phi/kernels/cpu/nms_kernel.cc
+++ b/paddle/phi/kernels/cpu/nms_kernel.cc
@@ -16,16 +16,17 @@
 #include "paddle/phi/backends/cpu/cpu_context.h"
 
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
 #include "paddle/phi/kernels/funcs/diagonal.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 
 namespace phi {
 
 template <typename T>
-static void NMS(const T* boxes_data,
-                int64_t* output_data,
-                float threshold,
-                int64_t num_boxes) {
+static int64_t NMS(const T* boxes_data,
+                   int64_t* output_data,
+                   float threshold,
+                   int64_t num_boxes) {
   auto num_masks = CeilDivide(num_boxes, 64);
   std::vector<uint64_t> masks(num_masks, 0);
 
@@ -54,9 +55,13 @@ static void NMS(const T* boxes_data,
     output_data[output_data_idx++] = i;
   }
 
+  int64_t num_keep_boxes = output_data_idx;
+
   for (; output_data_idx < num_boxes; ++output_data_idx) {
     output_data[output_data_idx] = 0;
   }
+
+  return num_keep_boxes;
 }
 
 template <typename T, typename Context>
@@ -64,8 +69,15 @@ void NMSKernel(const Context& dev_ctx,
                const DenseTensor& boxes,
                float threshold,
                DenseTensor* output) {
-  auto output_data = dev_ctx.template Alloc<int64_t>(output);
-  NMS<T>(boxes.data<T>(), output_data, threshold, boxes.dims()[0]);
+  int64_t num_boxes = boxes.dims()[0];
+  DenseTensor output_tmp;
+  output_tmp.Resize(phi::make_ddim({num_boxes}));
+  auto output_tmp_data = dev_ctx.template Alloc<int64_t>(&output_tmp);
+
+  int64_t num_keep_boxes =
+      NMS<T>(boxes.data<T>(), output_tmp_data, threshold, num_boxes);
+  auto slice_out = output_tmp.Slice(0, num_keep_boxes);
+  phi::Copy(dev_ctx, slice_out, dev_ctx.GetPlace(), false, output);
 }
 
 }  // namespace phi
diff --git a/paddle/phi/kernels/gpu/nms_kernel.cu b/paddle/phi/kernels/gpu/nms_kernel.cu
index 490753f1313..dcc6d6e2b45 100644
--- a/paddle/phi/kernels/gpu/nms_kernel.cu
+++ b/paddle/phi/kernels/gpu/nms_kernel.cu
@@ -59,7 +59,6 @@ void NMSKernel(const Context& dev_ctx,
                const DenseTensor& boxes,
                float threshold,
                DenseTensor* output) {
-  auto* output_data = dev_ctx.template Alloc<int64_t>(output);
   const int64_t num_boxes = boxes.dims()[0];
   const auto blocks_per_line = CeilDivide(num_boxes, threadsPerBlock);
   dim3 block(threadsPerBlock);
@@ -93,11 +92,13 @@ void NMSKernel(const Context& dev_ctx,
       }
     }
   }
+  output->Resize(phi::make_ddim({last_box_num}));
+  auto* output_data = dev_ctx.template Alloc<int64_t>(output);
   paddle::memory::Copy(dev_ctx.GetPlace(),
                        output_data,
                        phi::CPUPlace(),
                        output_host,
-                       sizeof(int64_t) * num_boxes,
+                       sizeof(int64_t) * last_box_num,
                        dev_ctx.stream());
 }
 }  // namespace phi
diff --git a/python/paddle/fluid/tests/unittests/test_nms_op.py b/python/paddle/fluid/tests/unittests/test_nms_op.py
index a81a46e1140..cbd24d4ddf2 100755
--- a/python/paddle/fluid/tests/unittests/test_nms_op.py
+++ b/python/paddle/fluid/tests/unittests/test_nms_op.py
@@ -65,7 +65,7 @@ def nms(boxes, nms_threshold):
             else:
                 continue
 
-    return selected_indices
+    return selected_indices[:cnt]
 
 
 class TestNMSOp(OpTest):
diff --git a/python/paddle/vision/ops.py b/python/paddle/vision/ops.py
index a83647349d9..2a17036693f 100755
--- a/python/paddle/vision/ops.py
+++ b/python/paddle/vision/ops.py
@@ -1611,7 +1611,9 @@ def nms(boxes,
     import paddle
     if category_idxs is None:
         sorted_global_indices = paddle.argsort(scores, descending=True)
-        return _nms(boxes[sorted_global_indices], iou_threshold)
+        sorted_keep_boxes_indices = _nms(boxes[sorted_global_indices],
+                                         iou_threshold)
+        return sorted_global_indices[sorted_keep_boxes_indices]
 
     if top_k is not None:
         assert top_k <= scores.shape[
-- 
GitLab