From 877bb24389f2adafa78d9eeb948bd8b5221bcd5b Mon Sep 17 00:00:00 2001 From: HaiyangP <46739135+HaiyangPeng@users.noreply.github.com> Date: Wed, 19 Apr 2023 17:00:14 +0800 Subject: [PATCH] Optional nms methods supporting for rcnn (#1274) * Support to output the instance segmentation results on an original input image, rather than on a padded one. * modifications according to freedenS's suggestions * modification for padding * merge preprocessImg and preprocessImg_ and delete all the debug codes. * bug of reference passing * make the vars consistent with the previous rcnn.cpp and update the README * support different nms methods * cxx11 * add contributors and delete interpretations * update README * fix some problems * fix the bugs --------- Co-authored-by: benkang wang --- rcnn/BatchedNms.cu | 39 +++++++++++++++++++++++++++++++++------ rcnn/BatchedNmsPlugin.h | 25 +++++++++++++++---------- rcnn/README.md | 4 +++- rcnn/rcnn.cpp | 23 ++++++++++++++++++++++- 4 files changed, 73 insertions(+), 18 deletions(-) diff --git a/rcnn/BatchedNms.cu b/rcnn/BatchedNms.cu index ec12c6d..96aac52 100755 --- a/rcnn/BatchedNms.cu +++ b/rcnn/BatchedNms.cu @@ -26,7 +26,7 @@ namespace cub = thrust::cuda_cub::cub; namespace nvinfer1 { __global__ void batched_nms_kernel( - const float threshold, const int num_detections, + const int nms_method, const float threshold, const int num_detections, const int *indices, float *scores, const float *classes, const float4 *boxes) { // Go through detections by descending score @@ -50,18 +50,44 @@ __global__ void batched_nms_kernel( float marea = (mbox.z - mbox.x) * (mbox.w - mbox.y); float inter = w * h; float overlap = inter / (iarea + marea - inter); - if (overlap > threshold) { - scores[i] = 0.0f; + float sigma = 0.5; // this is an empirical value + // printf("nms_method: %d", nms_method); + //nms methods selection in the second stage + // 0: original nms + // 1: soft-nms (linear) + // 2: soft-nms (gaussian) + // printf("nms_method: ", nms_method); + switch (nms_method) + { + case 0: + if (overlap > threshold) { + scores[i] = 0.0f; + } + break; + case 1: + if (overlap > threshold) { + scores[i] = (1 - overlap) * scores[i]; + } + break; + case 2: + if (overlap > threshold) { + scores[i] = std::exp(-(overlap * overlap) / sigma) * scores[i]; + } + break; + default: + if (overlap > threshold) { + scores[i] = 0.0f; + } + break; } } } - // Sync discarded detections __syncthreads(); } } -int batchedNms(int batch_size, +int batchedNms(int nms_method, int batch_size, const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs, size_t count, int detections_per_im, float nms_thresh, void *workspace, size_t workspace_size, cudaStream_t stream) { @@ -111,8 +137,9 @@ int batchedNms(int batch_size, // Launch actual NMS kernel - 1 block with each thread handling n detections // TODO: different device has differnet max threads const int max_threads = 1024; + int num_per_thread = ceil(static_cast(num_detections) / max_threads); - batched_nms_kernel << > > (nms_thresh, num_detections, + batched_nms_kernel << > > (nms_method, nms_thresh, num_detections, indices_sorted, scores_sorted, in_classes, in_boxes); // Re-sort with updated scores diff --git a/rcnn/BatchedNmsPlugin.h b/rcnn/BatchedNmsPlugin.h index 1d97627..7607832 100755 --- a/rcnn/BatchedNmsPlugin.h +++ b/rcnn/BatchedNmsPlugin.h @@ -13,7 +13,7 @@ using namespace nvinfer1; #define PLUGIN_NAMESPACE "" namespace nvinfer1 { -int batchedNms(int batchSize, +int batchedNms(int nms_method, int batchSize, const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs, size_t count, int detections_per_im, float nms_thresh, void *workspace, size_t workspace_size, cudaStream_t stream); @@ -28,6 +28,7 @@ int batchedNms(int batchSize, Description: implement batched nms */ class BatchedNmsPlugin : public IPluginV2Ext { + int _nms_method; float _nms_thresh; int _detections_per_im; @@ -36,32 +37,36 @@ class BatchedNmsPlugin : public IPluginV2Ext { protected: void deserialize(void const* data, size_t length) { const char* d = static_cast(data); + read(d, _nms_method); read(d, _nms_thresh); read(d, _detections_per_im); read(d, _count); } - size_t getSerializationSize() const TRT_NOEXCEPT override { - return sizeof(_nms_thresh) + sizeof(_detections_per_im) + size_t getSerializationSize() const override { + return sizeof(_nms_method) + sizeof(_nms_thresh) + sizeof(_detections_per_im) + sizeof(_count); } void serialize(void *buffer) const TRT_NOEXCEPT override { char* d = static_cast(buffer); + write(d, _nms_method); write(d, _nms_thresh); write(d, _detections_per_im); write(d, _count); } public: - BatchedNmsPlugin(float nms_thresh, int detections_per_im) - : _nms_thresh(nms_thresh), _detections_per_im(detections_per_im) { + BatchedNmsPlugin(int nms_method, float nms_thresh, int detections_per_im) + : _nms_method(nms_method), _nms_thresh(nms_thresh), _detections_per_im(detections_per_im) { + assert(nms_method >= 0); assert(nms_thresh > 0); assert(detections_per_im > 0); } - BatchedNmsPlugin(float nms_thresh, int detections_per_im, size_t count) - : _nms_thresh(nms_thresh), _detections_per_im(detections_per_im), _count(count) { + BatchedNmsPlugin(int nms_method, float nms_thresh, int detections_per_im, size_t count) + : _nms_method(nms_method), _nms_thresh(nms_thresh), _detections_per_im(detections_per_im), _count(count) { + assert(nms_method >= 0); assert(nms_thresh > 0); assert(detections_per_im > 0); assert(count > 0); @@ -101,7 +106,7 @@ class BatchedNmsPlugin : public IPluginV2Ext { size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { static int size = -1; if (size < 0) { - size = batchedNms(maxBatchSize, nullptr, nullptr, _count, + size = batchedNms(_nms_method, maxBatchSize, nullptr, nullptr, _count, _detections_per_im, _nms_thresh, nullptr, 0, nullptr); } @@ -111,7 +116,7 @@ class BatchedNmsPlugin : public IPluginV2Ext { int enqueue(int batchSize, const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs, void *workspace, cudaStream_t stream) TRT_NOEXCEPT override { - return batchedNms(batchSize, inputs, outputs, _count, + return batchedNms(_nms_method, batchSize, inputs, outputs, _count, _detections_per_im, _nms_thresh, workspace, getWorkspaceSize(batchSize), stream); } @@ -152,7 +157,7 @@ class BatchedNmsPlugin : public IPluginV2Ext { } IPluginV2Ext *clone() const TRT_NOEXCEPT override { - return new BatchedNmsPlugin(_nms_thresh, _detections_per_im, _count); + return new BatchedNmsPlugin(_nms_method, _nms_thresh, _detections_per_im, _count); } private: diff --git a/rcnn/README.md b/rcnn/README.md index 2f0ae52..e84bedf 100755 --- a/rcnn/README.md +++ b/rcnn/README.md @@ -1,6 +1,6 @@ # Rcnn -The Pytorch implementation is [facebookresearch/detectron2](https://github.com/facebookresearch/detectron2). Now, outputting instance segmentation results on the original image size is available, which is more convenient for engineering applications. +The Pytorch implementation is [facebookresearch/detectron2](https://github.com/facebookresearch/detectron2). Now, outputting instance segmentation results on the original image size and selecting different nms methods are available, which is more convenient for engineering applications. ## Models @@ -21,7 +21,9 @@ TensorRT 8.x is supported and you can use it. ## Contributors + + ## How to Run diff --git a/rcnn/rcnn.cpp b/rcnn/rcnn.cpp index 6923887..26265e1 100755 --- a/rcnn/rcnn.cpp +++ b/rcnn/rcnn.cpp @@ -52,6 +52,13 @@ static const char* INPUT_NODE_NAME = "images"; static const std::vector OUTPUT_NAMES = { "scores", "boxes", "labels", "masks" }; +//nms methods selection in the second stage +// 0: original nms +// 1: soft-nms (linear) +// 2: soft-nms (gaussian) +static int NMS_METHOD = 1; +static std::vector NMS_METHOD_VEC = {0, 1, 2}; + std::vector GenerateAnchors(const std::vector& anchor_sizes, const std::vector& aspect_ratios) { std::vector res; @@ -185,7 +192,7 @@ void BoxHead(INetworkDefinition *network, std::map& weight // nms std::vector nmsInput = { predictorDecodeLayer->getOutput(0), predictorDecodeLayer->getOutput(1), predictorDecodeLayer->getOutput(2) }; - auto batchedNmsPlugin = BatchedNmsPlugin(NMS_THRESH_TEST, DETECTIONS_PER_IMAGE); + auto batchedNmsPlugin = BatchedNmsPlugin(NMS_METHOD, NMS_THRESH_TEST, DETECTIONS_PER_IMAGE); auto batchedNmsLayer = network->addPluginV2(nmsInput.data(), nmsInput.size(), batchedNmsPlugin); // instances @@ -375,6 +382,20 @@ bool parse_args(int argc, char** argv, std::string& wtsFile, std::string& engine } int main(int argc, char** argv) { + + int flag = 0; + for (int &item : NMS_METHOD_VEC) { + if (item == NMS_METHOD) { + flag = 1; + printf("The nms method %d is applied.\n", NMS_METHOD); + break; + } + } + if (flag == 0) { + printf("[WARNING] The nms_method %d is not supported, please choose from [0, 1, 2].\n", NMS_METHOD); + printf("[WARNING] To make the nms robust, the default nms method 0 is applied.\n"); + NMS_METHOD = 0; + } // calculate size calculateSize();