diff --git a/paddle/fluid/operators/detection/matrix_nms_op.cc b/paddle/fluid/operators/detection/matrix_nms_op.cc index f7d45bc85bf6b14f07e9bfda3615e0b2d51a09f1..713c2dc7fe9c15916fba15dcfffe6c5f7e2a3958 100644 --- a/paddle/fluid/operators/detection/matrix_nms_op.cc +++ b/paddle/fluid/operators/detection/matrix_nms_op.cc @@ -12,6 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. limitations under the License. */ #include "paddle/fluid/framework/op_registry.h" +#include "paddle/fluid/framework/op_version_registry.h" #include "paddle/fluid/operators/detection/nms_util.h" namespace paddle { @@ -59,6 +60,9 @@ class MatrixNMSOp : public framework::OperatorWithKernel { } ctx->SetOutputDim("Out", {box_dims[1], box_dims[2] + 2}); ctx->SetOutputDim("Index", {box_dims[1], 1}); + if (ctx->HasOutput("RoisNum")) { + ctx->SetOutputDim("RoisNum", {-1}); + } if (!ctx->IsRuntime()) { ctx->SetLoDLevel("Out", std::max(ctx->GetLoDLevel("BBoxes"), 1)); ctx->SetLoDLevel("Index", std::max(ctx->GetLoDLevel("BBoxes"), 1)); @@ -259,8 +263,10 @@ class MatrixNMSKernel : public framework::OpKernel { std::vector offsets = {0}; std::vector detections; std::vector indices; + std::vector num_per_batch; detections.reserve(out_dim * num_boxes * batch_size); indices.reserve(num_boxes * batch_size); + num_per_batch.reserve(batch_size); for (int i = 0; i < batch_size; ++i) { scores_slice = scores->Slice(i, i + 1); scores_slice.Resize({score_dims[1], score_dims[2]}); @@ -272,6 +278,7 @@ class MatrixNMSKernel : public framework::OpKernel { background_label, nms_top_k, keep_top_k, normalized, score_threshold, post_threshold, use_gaussian, gaussian_sigma); offsets.push_back(offsets.back() + num_out); + num_per_batch.emplace_back(num_out); } int64_t num_kept = offsets.back(); @@ -285,6 +292,12 @@ class MatrixNMSKernel : public framework::OpKernel { std::copy(indices.begin(), indices.end(), index->data()); } + if (ctx.HasOutput("RoisNum")) { + auto* rois_num = ctx.Output("RoisNum"); + rois_num->mutable_data({batch_size}, ctx.GetPlace()); + std::copy(num_per_batch.begin(), num_per_batch.end(), + rois_num->data()); + } framework::LoD lod; lod.emplace_back(offsets); outs->set_lod(lod); @@ -355,6 +368,8 @@ class MatrixNMSOpMaker : public framework::OpProtoAndCheckerMaker { "(LoDTensor) A 2-D LoDTensor with shape [No, 1] represents the " "index of selected bbox. The index is the absolute index cross " "batches."); + AddOutput("RoisNum", "(Tensor), Number of RoIs in each images.") + .AsDispensable(); AddComment(R"DOC( This operator does multi-class matrix non maximum suppression (NMS) on batched boxes and scores. @@ -369,7 +384,9 @@ This operator support multi-class and batched inputs. It applying NMS independently for each class. The outputs is a 2-D LoDTenosr, for each image, the offsets in first dimension of LoDTensor are called LoD, the number of offset is N + 1, where N is the batch size. If LoD[i + 1] - LoD[i] == 0, -means there is no detected bbox for this image. +means there is no detected bbox for this image. Now this operator has one more +ouput, which is RoisNum. The size of RoisNum is N, RoisNum[i] means the number of +detected bbox for this image. For more information on Matrix NMS, please refer to: https://arxiv.org/abs/2003.10152 @@ -387,3 +404,8 @@ REGISTER_OPERATOR( paddle::framework::EmptyGradOpMaker); REGISTER_OP_CPU_KERNEL(matrix_nms, ops::MatrixNMSKernel, ops::MatrixNMSKernel); +REGISTER_OP_VERSION(matrix_nms) + .AddCheckpoint( + R"ROC(Upgrade matrix_nms: add a new output [RoisNum].)ROC", + paddle::framework::compatible::OpVersionDesc().NewOutput( + "RoisNum", "The number of RoIs in each image.")); diff --git a/paddle/fluid/pybind/op_function_generator.cc b/paddle/fluid/pybind/op_function_generator.cc index 8288f1852c27b1cf256e0c171bb4bf962e6caf6b..7f2736a9b1d4149f54c8eb516d1884141c90d446 100644 --- a/paddle/fluid/pybind/op_function_generator.cc +++ b/paddle/fluid/pybind/op_function_generator.cc @@ -74,6 +74,7 @@ std::map> op_outs_map = { {"unique", {"Out", "Index", "Indices", "Counts"}}, {"generate_proposals", {"RpnRois", "RpnRoiProbs", "RpnRoisNum"}}, {"collect_fpn_proposals", {"FpnRois", "RoisNum"}}, + {"matrix_nms", {"Out", "Index", "RoisNum"}}, {"distribute_fpn_proposals", {"MultiFpnRois", "RestoreIndex", "MultiLevelRoIsNum"}}, {"moving_average_abs_max_scale", {"OutScale", "OutAccum", "OutState"}}, diff --git a/python/paddle/fluid/tests/unittests/test_matrix_nms_op.py b/python/paddle/fluid/tests/unittests/test_matrix_nms_op.py index cf756ae8384486ea30ebc2fd25079484061ed380..2bbacc316f6e6095453a7875c78f6a0a4c677ebe 100644 --- a/python/paddle/fluid/tests/unittests/test_matrix_nms_op.py +++ b/python/paddle/fluid/tests/unittests/test_matrix_nms_op.py @@ -201,7 +201,8 @@ class TestMatrixNMSOp(OpTest): self.inputs = {'BBoxes': boxes, 'Scores': scores} self.outputs = { 'Out': (nmsed_outs, [lod]), - 'Index': (index_outs[:, None], [lod]) + 'Index': (index_outs[:, None], [lod]), + 'RoisNum': np.array(lod).astype('int32') } self.attrs = { 'background_label': 0,