haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

RCNN upgrade to support TensorRT 7&8 (#1256)

Browse Source 
* rcnn upgrade to support TensorRT 8.

* Update RpnDecodePlugin.h

Remove Chinese

* Update rcnn.cpp

Remove Chinese

* Update backbone.hpp

Remove Chinese

* Update backbone.hpp

* Update rcnn.cpp

* Update rcnn.cpp

* Update rcnn.cpp

* Update MaskRcnnInferencePlugin.h

* rcnn upgrade to support TensorRT 8.x

* rcnn upgrade to support TensorRT 8.x

* Update macros.h

* Update README.md

---------

Co-authored-by: nengwp <nengwp@github.ai>
Co-authored-by: Wang Xinyu <shaywxy@gmail.com>
This commit is contained in:
nengwp 2023-03-07 19:26:53 +08:00 committed by GitHub
parent ef22b1d03b
commit f2521a5ffa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
23 changed files with 286 additions and 210 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
README.md
View File

@ -14,7 +14,7 @@ The basic workflow of TensorRTx is:
4. Load the TensorRT engine and run inference.
## News
- `1 Mar 2023`. [Nengwp](https://github.com/nengwp): [RCNN](./rcnn) and [UNet](./unet) upgrade to support TensorRT 8.
- `18 Dec 2022`. [YOLOv5](./yolov5) upgrade to support v7.0, including instance segmentation.
- `12 Dec 2022`. [East-Face](https://github.com/East-Face): [UNet](./unet) upgrade to support v3.0 of [Pytorch-UNet](https://github.com/milesial/Pytorch-UNet).
- `26 Oct 2022`. [ausk](https://github.com/ausk): YoloP(You Only Look Once for Panopitic Driving Perception).
@ -29,7 +29,6 @@ The basic workflow of TensorRTx is:
- `19 Oct 2021`. [liuqi123123](https://github.com/liuqi123123) added cuda preprossing for yolov5, preprocessing + inference is 3x faster when batchsize=8.
- `18 Oct 2021`. [xupengao](https://github.com/xupengao): YOLOv5 updated to v6.0, supporting n/s/m/l/x/n6/s6/m6/l6/x6.
- `31 Aug 2021`. [FamousDirector](https://github.com/FamousDirector): update retinaface to support TensorRT 8.0.
- `27 Aug 2021`. [HaiyangPeng](https://github.com/HaiyangPeng): add a python wrapper for hrnet segmentation.
## Tutorials

20
rcnn/BatchedNms.cu Normal file → Executable file
View File

@ -3,7 +3,6 @@
#include <thrust/sequence.h>
#include <thrust/execution_policy.h>
#include <thrust/gather.h>
#include <thrust/system/cuda/detail/cub/device/device_radix_sort.cuh>
#include <cmath>
#include <algorithm>
#include <iostream>
@ -12,6 +11,17 @@
#include <vector>
#include "BatchedNmsPlugin.h"
#include "./cuda_utils.h"
#include "macros.h"
#ifdef CUDA_11
#include <cub/device/device_radix_sort.cuh>
#include <cub/iterator/counting_input_iterator.cuh>
#else
#include <thrust/system/cuda/detail/cub/device/device_radix_sort.cuh>
#include <thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh>
namespace cub = thrust::cuda_cub::cub;
#endif
namespace nvinfer1 {
@ -52,7 +62,7 @@ __global__ void batched_nms_kernel(
}
int batchedNms(int batch_size,
const void *const *inputs, void **outputs,
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) {
@ -63,7 +73,7 @@ int batchedNms(int batch_size,
workspace_size += get_size_aligned<float>(count); // scores_sorted
size_t temp_size_sort = 0;
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(
cub::DeviceRadixSort::SortPairsDescending(
static_cast<void*>(nullptr), temp_size_sort,
static_cast<float*>(nullptr),
static_cast<float*>(nullptr),
@ -95,7 +105,7 @@ int batchedNms(int batch_size,
// Sort scores and corresponding indices
int num_detections = count;
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
in_scores, scores_sorted, indices, indices_sorted, num_detections, 0, sizeof(*scores_sorted) * 8, stream);
// Launch actual NMS kernel - 1 block with each thread handling n detections
@ -106,7 +116,7 @@ int batchedNms(int batch_size,
indices_sorted, scores_sorted, in_classes, in_boxes);
// Re-sort with updated scores
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
scores_sorted, scores_sorted, indices_sorted, indices,
num_detections, 0, sizeof(*scores_sorted) * 8, stream);

57
rcnn/BatchedNmsPlugin.h Normal file → Executable file
View File

@ -4,6 +4,7 @@
#include <vector>
#include <cassert>
#include "macros.h"
using namespace nvinfer1;
@ -13,7 +14,7 @@ using namespace nvinfer1;
namespace nvinfer1 {
int batchedNms(int batchSize,
const void *const *inputs, void **outputs,
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);
@ -40,12 +41,12 @@ class BatchedNmsPlugin : public IPluginV2Ext {
read(d, _count);
}
size_t getSerializationSize() const override {
size_t getSerializationSize() const TRT_NOEXCEPT override {
return sizeof(_nms_thresh) + sizeof(_detections_per_im)
+ sizeof(_count);
}
void serialize(void *buffer) const override {
void serialize(void *buffer) const TRT_NOEXCEPT override {
char* d = static_cast<char*>(buffer);
write(d, _nms_thresh);
write(d, _detections_per_im);
@ -70,34 +71,34 @@ class BatchedNmsPlugin : public IPluginV2Ext {
this->deserialize(data, length);
}
const char *getPluginType() const override {
const char *getPluginType() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
int getNbOutputs() const override {
int getNbOutputs() const TRT_NOEXCEPT override {
return 3;
}
Dims getOutputDimensions(int index,
const Dims *inputs, int nbInputDims) override {
const Dims *inputs, int nbInputDims) TRT_NOEXCEPT override {
assert(nbInputDims == 3);
assert(index < this->getNbOutputs());
return Dims2(_detections_per_im, index == 1 ? 4 : 1);
}
bool supportsFormat(DataType type, PluginFormat format) const override {
bool supportsFormat(DataType type, PluginFormat format) const TRT_NOEXCEPT override {
return type == DataType::kFLOAT && format == PluginFormat::kLINEAR;
}
int initialize() override { return 0; }
int initialize() TRT_NOEXCEPT override { return 0; }
void terminate() override {}
void terminate() TRT_NOEXCEPT override {}
size_t getWorkspaceSize(int maxBatchSize) const override {
size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override {
static int size = -1;
if (size < 0) {
size = batchedNms(maxBatchSize, nullptr, nullptr, _count,
@ -108,40 +109,40 @@ class BatchedNmsPlugin : public IPluginV2Ext {
}
int enqueue(int batchSize,
const void *const *inputs, void **outputs,
void *workspace, cudaStream_t stream) override {
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
void *workspace, cudaStream_t stream) TRT_NOEXCEPT override {
return batchedNms(batchSize, inputs, outputs, _count,
_detections_per_im, _nms_thresh,
workspace, getWorkspaceSize(batchSize), stream);
}
void destroy() override {
void destroy() TRT_NOEXCEPT override {
delete this;
}
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
void setPluginNamespace(const char *N) override {
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {
}
// IPluginV2Ext Methods
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const {
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override {
assert(index < 3);
return DataType::kFLOAT;
}
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted,
int nbInputs) const {
int nbInputs) const TRT_NOEXCEPT override {
return false;
}
bool canBroadcastInputAcrossBatch(int inputIndex) const { return false; }
bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override { return false; }
void configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast,
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) {
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) TRT_NOEXCEPT override {
assert(*inputTypes == nvinfer1::DataType::kFLOAT &&
floatFormat == nvinfer1::PluginFormat::kLINEAR);
assert(nbInputs == 3);
@ -150,7 +151,7 @@ class BatchedNmsPlugin : public IPluginV2Ext {
_count = inputDims[0].d[0];
}
IPluginV2Ext *clone() const override {
IPluginV2Ext *clone() const TRT_NOEXCEPT override {
return new BatchedNmsPlugin(_nms_thresh, _detections_per_im, _count);
}
@ -170,24 +171,24 @@ class BatchedNmsPluginCreator : public IPluginCreator {
public:
BatchedNmsPluginCreator() {}
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
const char *getPluginName() const override {
const char *getPluginName() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override {
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) TRT_NOEXCEPT override {
return new BatchedNmsPlugin(serialData, serialLength);
}
void setPluginNamespace(const char *N) override {}
const PluginFieldCollection *getFieldNames() override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) override { return nullptr; }
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {}
const PluginFieldCollection *getFieldNames() TRT_NOEXCEPT override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) TRT_NOEXCEPT override { return nullptr; }
};
REGISTER_TENSORRT_PLUGIN(BatchedNmsPluginCreator);

16
rcnn/CMakeLists.txt Normal file → Executable file
View File

@ -1,11 +1,11 @@
cmake_minimum_required(VERSION 2.6)
cmake_minimum_required(VERSION 3.1)
project(rcnn)
add_definitions(-std=c++11)
add_definitions(-std=c++14)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_STANDARD 14)
set(CMAKE_BUILD_TYPE Debug)
set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};--extended-lambda)
@ -15,13 +15,13 @@ find_package(CUDA REQUIRED)
include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda-10.2/include)
link_directories(/usr/local/cuda-10.2/lib64)
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/home/jushi/TensorRT-7.2.1.6/include)
link_directories(/home/jushi/TensorRT-7.2.1.6/lib)
include_directories(/home/jushi/TensorRT-8.2.1.6/include)
link_directories(/home/jushi/TensorRT-8.2.1.6/lib)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
cuda_add_library(myplugins SHARED ${PROJECT_SOURCE_DIR}/BatchedNms.cu ${PROJECT_SOURCE_DIR}/PredictorDecode.cu ${PROJECT_SOURCE_DIR}/RoiAlign.cu ${PROJECT_SOURCE_DIR}/RpnDecode.cu ${PROJECT_SOURCE_DIR}/RpnNms.cu ${PROJECT_SOURCE_DIR}/MaskRcnnInference.cu)
target_link_libraries(myplugins nvinfer cudart)

3
rcnn/MaskRcnnInference.cu Normal file → Executable file
View File

@ -1,4 +1,5 @@
#include "MaskRcnnInferencePlugin.h"
#include "macros.h"
namespace nvinfer1 {
@ -31,7 +32,7 @@ __global__ void MaskRcnnInferenceKernel(
}
int maskRcnnInference(int batchSize,
const void *const *inputs, void **outputs,
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
int detections_per_im, int output_size, int num_classes, cudaStream_t stream) {
for (int batch = 0; batch < batchSize; batch++) {

57
rcnn/MaskRcnnInferencePlugin.h Normal file → Executable file
View File

@ -4,6 +4,7 @@
#include <vector>
#include <cassert>
#include "macros.h"
using namespace nvinfer1;
@ -13,7 +14,7 @@ using namespace nvinfer1;
namespace nvinfer1 {
int maskRcnnInference(int batchSize,
const void *const *inputs, void **outputs,
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
int detections_per_im, int output_size, int num_classes, cudaStream_t stream);
/*
input1: indices{C, 1} C->topk
@ -34,10 +35,10 @@ class MaskRcnnInferencePlugin : public IPluginV2Ext {
read(d, _output_size);
read(d, _num_classes);
}
size_t getSerializationSize() const override {
size_t getSerializationSize() const TRT_NOEXCEPT override {
return sizeof(_detections_per_im) + sizeof(_output_size) + sizeof(_num_classes);
}
void serialize(void *buffer) const override {
void serialize(void *buffer) const TRT_NOEXCEPT override {
char* d = static_cast<char*>(buffer);
write(d, _detections_per_im);
write(d, _output_size);
@ -59,55 +60,55 @@ class MaskRcnnInferencePlugin : public IPluginV2Ext {
MaskRcnnInferencePlugin(void const* data, size_t length) {
this->deserialize(data, length);
}
const char *getPluginType() const override {
const char *getPluginType() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
int getNbOutputs() const override {
int getNbOutputs() const TRT_NOEXCEPT override {
return 1;
}
Dims getOutputDimensions(int index,
const Dims *inputs, int nbInputDims) override {
const Dims *inputs, int nbInputDims) TRT_NOEXCEPT override {
assert(index < this->getNbOutputs());
return Dims4(_detections_per_im, 1, _output_size, _output_size);
}
bool supportsFormat(DataType type, PluginFormat format) const override {
bool supportsFormat(DataType type, PluginFormat format) const TRT_NOEXCEPT override {
return type == DataType::kFLOAT && format == PluginFormat::kLINEAR;
}
int initialize() override { return 0; }
void terminate() override {}
size_t getWorkspaceSize(int maxBatchSize) const override {
int initialize() TRT_NOEXCEPT override { return 0; }
void terminate() TRT_NOEXCEPT override {}
size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override {
return 0;
}
int enqueue(int batchSize,
const void *const *inputs, void **outputs,
void *workspace, cudaStream_t stream) override {
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
void *workspace, cudaStream_t stream) TRT_NOEXCEPT override {
return maskRcnnInference(batchSize, inputs, outputs,
_detections_per_im, _output_size, _num_classes, stream);
}
void destroy() override {
void destroy() TRT_NOEXCEPT override {
delete this;
}
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
void setPluginNamespace(const char *N) override {
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {
}
// IPluginV2Ext Methods
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const {
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override {
assert(index < 1);
return DataType::kFLOAT;
}
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted,
int nbInputs) const {
int nbInputs) const TRT_NOEXCEPT override {
return false;
}
bool canBroadcastInputAcrossBatch(int inputIndex) const { return false; }
bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override { return false; }
void configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast,
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) {
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) TRT_NOEXCEPT override {
assert(*inputTypes == nvinfer1::DataType::kFLOAT &&
floatFormat == nvinfer1::PluginFormat::kLINEAR);
assert(nbInputs == 2);
@ -117,7 +118,7 @@ class MaskRcnnInferencePlugin : public IPluginV2Ext {
assert(inputDims[1].d[3] == _output_size);
_num_classes = inputDims[1].d[1];
}
IPluginV2Ext *clone() const override {
IPluginV2Ext *clone() const TRT_NOEXCEPT override {
return new MaskRcnnInferencePlugin(_detections_per_im, _output_size, _num_classes);
}
@ -135,21 +136,21 @@ class MaskRcnnInferencePlugin : public IPluginV2Ext {
class MaskRcnnInferencePluginCreator : public IPluginCreator {
public:
MaskRcnnInferencePluginCreator() {}
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
const char *getPluginName() const override {
const char *getPluginName() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override {
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) TRT_NOEXCEPT override {
return new MaskRcnnInferencePlugin(serialData, serialLength);
}
void setPluginNamespace(const char *N) override {}
const PluginFieldCollection *getFieldNames() override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) override { return nullptr; }
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {}
const PluginFieldCollection *getFieldNames() TRT_NOEXCEPT override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) TRT_NOEXCEPT override { return nullptr; }
};
REGISTER_TENSORRT_PLUGIN(MaskRcnnInferencePluginCreator);

18
rcnn/PredictorDecode.cu Normal file → Executable file
View File

@ -2,19 +2,27 @@
#include <thrust/sequence.h>
#include <thrust/execution_policy.h>
#include <thrust/gather.h>
#include <thrust/system/cuda/detail/cub/device/device_radix_sort.cuh>
#include <thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh>
#include <algorithm>
#include <cstdint>
#include "PredictorDecodePlugin.h"
#include "./cuda_utils.h"
#include "macros.h"
#ifdef CUDA_11
#include <cub/device/device_radix_sort.cuh>
#include <cub/iterator/counting_input_iterator.cuh>
#else
#include <thrust/system/cuda/detail/cub/device/device_radix_sort.cuh>
#include <thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh>
namespace cub = thrust::cuda_cub::cub;
#endif
namespace nvinfer1 {
int predictorDecode(int batchSize, const void *const *inputs,
void **outputs, unsigned int num_boxes, unsigned int num_classes,
void *TRT_CONST_ENQUEUE*outputs, unsigned int num_boxes, unsigned int num_classes,
unsigned int image_height, unsigned int image_width,
const std::vector<float>& bbox_reg_weights, void *workspace,
size_t workspace_size, cudaStream_t stream) {
@ -28,7 +36,7 @@ size_t workspace_size, cudaStream_t stream) {
workspace_size += get_size_aligned<float>(scores_size); // scores_sorted
size_t temp_size_sort = 0;
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(
cub::DeviceRadixSort::SortPairsDescending(
static_cast<void*>(nullptr), temp_size_sort,
static_cast<float*>(nullptr),
static_cast<float*>(nullptr),
@ -64,7 +72,7 @@ size_t workspace_size, cudaStream_t stream) {
auto out_classes = static_cast<float *>(outputs[2]) + batch * num_boxes;
// Only keep top n scores
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
in_scores, scores_sorted, indices, indices_sorted, scores_size, 0, sizeof(*scores_sorted) * 8, stream);
// Gather boxes

57
rcnn/PredictorDecodePlugin.h Normal file → Executable file
View File

@ -4,6 +4,7 @@
#include <cassert>
#include <vector>
#include "macros.h"
using namespace nvinfer1;
@ -14,7 +15,7 @@ using namespace nvinfer1;
namespace nvinfer1 {
int predictorDecode(int batchSize,
const void *const *inputs, void **outputs, unsigned int num_boxes,
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs, unsigned int num_boxes,
unsigned int num_classes, unsigned int image_height,
unsigned int image_width, const std::vector<float>& bbox_reg_weights,
void *workspace, size_t workspace_size, cudaStream_t stream);
@ -52,13 +53,13 @@ class PredictorDecodePlugin : public IPluginV2Ext {
}
}
size_t getSerializationSize() const override {
size_t getSerializationSize() const TRT_NOEXCEPT override {
return sizeof(_num_boxes) + sizeof(_num_classes) +
sizeof(_image_height) + sizeof(_image_width) + sizeof(size_t) +
sizeof(float)*_bbox_reg_weights.size();
}
void serialize(void *buffer) const override {
void serialize(void *buffer) const TRT_NOEXCEPT override {
char* d = static_cast<char*>(buffer);
write(d, _num_boxes);
write(d, _num_classes);
@ -87,34 +88,34 @@ class PredictorDecodePlugin : public IPluginV2Ext {
this->deserialize(data, length);
}
const char *getPluginType() const override {
const char *getPluginType() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
int getNbOutputs() const override {
int getNbOutputs() const TRT_NOEXCEPT override {
return 3;
}
Dims getOutputDimensions(int index,
const Dims *inputs, int nbInputDims) override {
const Dims *inputs, int nbInputDims) TRT_NOEXCEPT override {
assert(nbInputDims == 3);
assert(index < this->getNbOutputs());
return Dims2(_num_boxes, (index == 1 ? 4 : 1));
}
bool supportsFormat(DataType type, PluginFormat format) const override {
bool supportsFormat(DataType type, PluginFormat format) const TRT_NOEXCEPT override {
return type == DataType::kFLOAT && format == PluginFormat::kLINEAR;
}
int initialize() override { return 0; }
int initialize() TRT_NOEXCEPT override { return 0; }
void terminate() override {}
void terminate() TRT_NOEXCEPT override {}
size_t getWorkspaceSize(int maxBatchSize) const override {
size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override {
if (size < 0) {
size = predictorDecode(maxBatchSize, nullptr, nullptr,
_num_boxes, _num_classes, _image_height, _image_width,
@ -124,39 +125,39 @@ class PredictorDecodePlugin : public IPluginV2Ext {
}
int enqueue(int batchSize,
const void *const *inputs, void **outputs,
void *workspace, cudaStream_t stream) override {
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
void *workspace, cudaStream_t stream) TRT_NOEXCEPT override {
return predictorDecode(batchSize, inputs, outputs, _num_boxes,
_num_classes, _image_height, _image_width, _bbox_reg_weights,
workspace, getWorkspaceSize(batchSize), stream);
}
void destroy() override {
void destroy() TRT_NOEXCEPT override {
delete this;
};
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
void setPluginNamespace(const char *N) override {}
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {}
// IPluginV2Ext Methods
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const {
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override {
assert(index < this->getNbOutputs());
return DataType::kFLOAT;
}
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted,
int nbInputs) const {
int nbInputs) const TRT_NOEXCEPT override {
return false;
}
bool canBroadcastInputAcrossBatch(int inputIndex) const { return false; }
bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override { return false; }
void configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast,
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) {
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) TRT_NOEXCEPT override {
assert(*inputTypes == nvinfer1::DataType::kFLOAT &&
floatFormat == nvinfer1::PluginFormat::kLINEAR);
assert(nbInputs == 3);
@ -172,7 +173,7 @@ class PredictorDecodePlugin : public IPluginV2Ext {
_num_classes = scores_dims.d[1];
}
IPluginV2Ext *clone() const override {
IPluginV2Ext *clone() const TRT_NOEXCEPT override {
return new PredictorDecodePlugin(_num_boxes, _num_classes, _image_height, _image_width, _bbox_reg_weights);
}
@ -192,25 +193,25 @@ class PredictorDecodePluginCreator : public IPluginCreator {
public:
PredictorDecodePluginCreator() {}
const char *getPluginName() const override {
const char *getPluginName() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override {
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) TRT_NOEXCEPT override {
return new PredictorDecodePlugin(serialData, serialLength);
}
void setPluginNamespace(const char *N) override {}
const PluginFieldCollection *getFieldNames() override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) override { return nullptr; }
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {}
const PluginFieldCollection *getFieldNames() TRT_NOEXCEPT override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) TRT_NOEXCEPT override { return nullptr; }
};
REGISTER_TENSORRT_PLUGIN(PredictorDecodePluginCreator);

6
rcnn/README.md Normal file → Executable file
View File

@ -10,13 +10,19 @@ The Pytorch implementation is [facebookresearch/detectron2](https://github.com/f
## Test Environment
- GTX3090 / Ubuntu20.04 / cuda11 / cudnn8.0.4 / TensorRT8.1.1 / OpenCV4.5 form docker hakuyyf/tensorrtx:trt8_cuda11
- GTX2080Ti / Ubuntu16.04 / cuda10.2 / cudnn8.0.4 / TensorRT7.2.1 / OpenCV4.2
- GTX2080Ti / win10 / cuda10.2 / cudnn8.0.4 / TensorRT7.2.1 / OpenCV4.2 / VS2017 (need to replace function corresponding to the dirent.h and add "--extended-lambda" in CUDA C/C++ -> Command Line -> Other options)
TensorRT7.2 is recomended because Resize layer in 7.0 with kLINEAR mode is a little different with opencv. You can also implement data preprocess out of tensorrt if you want to use TensorRT7.0 or more previous version.
TensorRT 8.x is supported and you can use it.
**The result under fp32 is same to pytorch about 4 decimal places**!
## Contributors
<a href="https://github.com/nengwp"><img src="https://avatars.githubusercontent.com/u/44516353?s=96&v=4" width="40px;" alt=""/></a>
## How to Run
1. generate .wts from pytorch with .pkl or .pth

14
rcnn/RoiAlign.cu Normal file → Executable file
View File

@ -14,6 +14,16 @@
#include "RoiAlignPlugin.h"
#include "./cuda_utils.h"
#include "macros.h"
#ifdef CUDA_11
#include <cub/device/device_radix_sort.cuh>
#include <cub/iterator/counting_input_iterator.cuh>
#else
#include <thrust/system/cuda/detail/cub/device/device_radix_sort.cuh>
#include <thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh>
namespace cub = thrust::cuda_cub::cub;
#endif
namespace nvinfer1 {
template <typename T>
@ -118,7 +128,7 @@ __global__ void RoIAlignForward(
const float count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4
float output_val = 0.f;
bool max_flag = false;
// bool max_flag = false;
// e.g., iy = 0, 1
for (int iy = 0; iy < roi_bin_grid_h; iy++) {
const float y = roi_start_h + ph * bin_size_h +
@ -142,7 +152,7 @@ __global__ void RoIAlignForward(
}
}
int roiAlign(int batchSize, const void *const *inputs, void **outputs, int pooler_resolution, float spatial_scale,
int roiAlign(int batchSize, const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs, int pooler_resolution, float spatial_scale,
int sampling_ratio, int num_proposals, int out_channels, int feature_h, int feature_w, cudaStream_t stream) {
for (int batch = 0; batch < batchSize; batch++) {
auto in_boxes = static_cast<const float4 *>(inputs[0]) + batch * num_proposals;

57
rcnn/RoiAlignPlugin.h Normal file → Executable file
View File

@ -4,6 +4,7 @@
#include <cassert>
#include <vector>
#include "macros.h"
using namespace nvinfer1;
@ -12,7 +13,7 @@ using namespace nvinfer1;
#define PLUGIN_NAMESPACE ""
namespace nvinfer1 {
int roiAlign(int batchSize, const void *const *inputs, void **outputs,
int roiAlign(int batchSize, const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
int pooler_resolution, float spatial_scale, int sampling_ratio,
int num_proposals, int out_channels, int feature_h, int feature_w,
cudaStream_t stream);
@ -44,12 +45,12 @@ class RoiAlignPlugin : public IPluginV2Ext {
read(d, _feature_w);
}
size_t getSerializationSize() const override {
size_t getSerializationSize() const TRT_NOEXCEPT override {
return sizeof(_pooler_resolution) + sizeof(_spatial_scale) + sizeof(_sampling_ratio) +
sizeof(_num_proposals) + sizeof(_out_channels) + sizeof(_feature_h) + sizeof(_feature_w);
}
void serialize(void *buffer) const override {
void serialize(void *buffer) const TRT_NOEXCEPT override {
char* d = static_cast<char*>(buffer);
write(d, _pooler_resolution);
write(d, _spatial_scale);
@ -75,70 +76,70 @@ class RoiAlignPlugin : public IPluginV2Ext {
this->deserialize(data, length);
}
const char *getPluginType() const override {
const char *getPluginType() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
int getNbOutputs() const override {
int getNbOutputs() const TRT_NOEXCEPT override {
return 1;
}
Dims getOutputDimensions(int index,
const Dims *inputs, int nbInputDims) override {
const Dims *inputs, int nbInputDims) TRT_NOEXCEPT override {
assert(index < this->getNbOutputs());
return Dims4(_num_proposals, _out_channels, _pooler_resolution, _pooler_resolution);
}
bool supportsFormat(DataType type, PluginFormat format) const override {
bool supportsFormat(DataType type, PluginFormat format) const TRT_NOEXCEPT override {
return type == DataType::kFLOAT && format == PluginFormat::kLINEAR;
}
int initialize() override { return 0; }
int initialize() TRT_NOEXCEPT override { return 0; }
void terminate() override {}
void terminate() TRT_NOEXCEPT override {}
size_t getWorkspaceSize(int maxBatchSize) const override {
size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override {
return 0;
}
int enqueue(int batchSize,
const void *const *inputs, void **outputs,
void *workspace, cudaStream_t stream) override {
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
void *workspace, cudaStream_t stream) TRT_NOEXCEPT override {
return roiAlign(batchSize, inputs, outputs, _pooler_resolution, _spatial_scale, _sampling_ratio,
_num_proposals, _out_channels, _feature_h, _feature_w, stream);
}
void destroy() override {
void destroy() TRT_NOEXCEPT override {
delete this;
};
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
void setPluginNamespace(const char *N) override {
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {
}
// IPluginV2Ext Methods
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const {
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override {
assert(index < this->getNbOutputs());
return DataType::kFLOAT;
}
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted,
int nbInputs) const {
int nbInputs) const TRT_NOEXCEPT override {
return false;
}
bool canBroadcastInputAcrossBatch(int inputIndex) const { return false; }
bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override { return false; }
void configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast,
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) {
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) TRT_NOEXCEPT override {
assert(*inputTypes == nvinfer1::DataType::kFLOAT &&
floatFormat == nvinfer1::PluginFormat::kLINEAR);
assert(nbInputs == 2);
@ -151,7 +152,7 @@ class RoiAlignPlugin : public IPluginV2Ext {
_feature_w = feature_dims.d[2];
}
IPluginV2Ext *clone() const override {
IPluginV2Ext *clone() const TRT_NOEXCEPT override {
return new RoiAlignPlugin(_pooler_resolution, _spatial_scale, _sampling_ratio, _num_proposals,
_out_channels, _feature_h, _feature_w);
}
@ -172,25 +173,25 @@ class RoiAlignPluginCreator : public IPluginCreator {
public:
RoiAlignPluginCreator() {}
const char *getPluginName() const override {
const char *getPluginName() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override {
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) TRT_NOEXCEPT override {
return new RoiAlignPlugin(serialData, serialLength);
}
void setPluginNamespace(const char *N) override {}
const PluginFieldCollection *getFieldNames() override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) override { return nullptr; }
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {}
const PluginFieldCollection *getFieldNames() TRT_NOEXCEPT override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) TRT_NOEXCEPT override { return nullptr; }
};
REGISTER_TENSORRT_PLUGIN(RoiAlignPluginCreator);

17
rcnn/RpnDecode.cu Normal file → Executable file
View File

@ -5,18 +5,27 @@
#include <thrust/tabulate.h>
#include <thrust/count.h>
#include <thrust/find.h>
#include <thrust/system/cuda/detail/cub/device/device_radix_sort.cuh>
#include <algorithm>
#include <cstdint>
#include "RpnDecodePlugin.h"
#include "./cuda_utils.h"
#include "macros.h"
#ifdef CUDA_11
#include <cub/device/device_radix_sort.cuh>
#include <cub/iterator/counting_input_iterator.cuh>
#else
#include <thrust/system/cuda/detail/cub/device/device_radix_sort.cuh>
#include <thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh>
namespace cub = thrust::cuda_cub::cub;
#endif
namespace nvinfer1 {
int rpnDecode(int batch_size,
const void *const *inputs, void **outputs,
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
size_t height, size_t width, size_t image_height, size_t image_width, float stride,
const std::vector<float> &anchors, int top_n,
void *workspace, size_t workspace_size, cudaStream_t stream) {
@ -33,7 +42,7 @@ int rpnDecode(int batch_size,
size_t temp_size_sort = 0;
if (scores_size > top_n) {
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(
cub::DeviceRadixSort::SortPairsDescending(
static_cast<void*>(nullptr), temp_size_sort,
static_cast<float*>(nullptr),
static_cast<float*>(nullptr),
@ -70,7 +79,7 @@ int rpnDecode(int batch_size,
int num_detections = scores_size;
auto indices_filtered = indices;
if (num_detections > top_n) {
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
in_scores, scores_sorted, indices, indices_sorted, scores_size, 0, sizeof(*scores_sorted) * 8, stream);
indices_filtered = indices_sorted;
num_detections = top_n;

57
rcnn/RpnDecodePlugin.h Normal file → Executable file
View File

@ -4,6 +4,7 @@
#include <cassert>
#include <vector>
#include "macros.h"
using namespace nvinfer1;
@ -14,7 +15,7 @@ using namespace nvinfer1;
namespace nvinfer1 {
int rpnDecode(int batchSize, const void *const *inputs,
void **outputs, size_t height, size_t width, size_t image_height,
void *TRT_CONST_ENQUEUE*outputs, size_t height, size_t width, size_t image_height,
size_t image_width, float stride, const std::vector<float> &anchors,
int top_n, void *workspace, size_t workspace_size, cudaStream_t stream);
@ -54,13 +55,13 @@ class RpnDecodePlugin : public IPluginV2Ext {
read(d, _image_width);
}
size_t getSerializationSize() const override {
size_t getSerializationSize() const TRT_NOEXCEPT override {
return sizeof(_top_n)
+ sizeof(size_t) + sizeof(float) * _anchors.size() + sizeof(_stride)
+ sizeof(_height) + sizeof(_width) + sizeof(_image_height) + sizeof(_image_width);
}
void serialize(void *buffer) const override {
void serialize(void *buffer) const TRT_NOEXCEPT override {
char* d = static_cast<char*>(buffer);
write(d, _top_n);
write(d, _anchors.size());
@ -87,34 +88,34 @@ class RpnDecodePlugin : public IPluginV2Ext {
this->deserialize(data, length);
}
const char *getPluginType() const override {
const char *getPluginType() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
int getNbOutputs() const override {
int getNbOutputs() const TRT_NOEXCEPT override {
return 2;
}
Dims getOutputDimensions(int index,
const Dims *inputs, int nbInputDims) override {
const Dims *inputs, int nbInputDims) TRT_NOEXCEPT override {
assert(nbInputDims == 2);
assert(index < this->getNbOutputs());
return Dims2(_top_n, (index == 1 ? 4 : 1));
}
bool supportsFormat(DataType type, PluginFormat format) const override {
bool supportsFormat(DataType type, PluginFormat format) const TRT_NOEXCEPT override {
return type == DataType::kFLOAT && format == PluginFormat::kLINEAR;
}
int initialize() override { return 0; }
int initialize() TRT_NOEXCEPT override { return 0; }
void terminate() override {}
void terminate() TRT_NOEXCEPT override {}
size_t getWorkspaceSize(int maxBatchSize) const override {
size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override {
if (size < 0) {
size = rpnDecode(maxBatchSize, nullptr, nullptr, _height, _width, _image_height, _image_width, _stride,
_anchors, _top_n,
@ -124,39 +125,39 @@ class RpnDecodePlugin : public IPluginV2Ext {
}
int enqueue(int batchSize,
const void *const *inputs, void **outputs,
void *workspace, cudaStream_t stream) override {
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
void *workspace, cudaStream_t stream) TRT_NOEXCEPT override {
return rpnDecode(batchSize, inputs, outputs, _height, _width, _image_height, _image_width, _stride,
_anchors, _top_n, workspace, getWorkspaceSize(batchSize), stream);
}
void destroy() override {
void destroy() TRT_NOEXCEPT override {
delete this;
};
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
void setPluginNamespace(const char *N) override {
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {
}
// IPluginV2Ext Methods
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const {
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override {
assert(index < 3);
return DataType::kFLOAT;
}
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted,
int nbInputs) const {
int nbInputs) const TRT_NOEXCEPT override {
return false;
}
bool canBroadcastInputAcrossBatch(int inputIndex) const { return false; }
bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override { return false; }
void configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast,
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) {
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) TRT_NOEXCEPT override {
assert(*inputTypes == nvinfer1::DataType::kFLOAT &&
floatFormat == nvinfer1::PluginFormat::kLINEAR);
assert(nbInputs == 2);
@ -169,7 +170,7 @@ class RpnDecodePlugin : public IPluginV2Ext {
_width = scores_dims.d[2];
}
IPluginV2Ext *clone() const override {
IPluginV2Ext *clone() const TRT_NOEXCEPT override {
return new RpnDecodePlugin(_top_n, _anchors, _stride, _height, _width, _image_height, _image_width);
}
@ -189,25 +190,25 @@ class RpnDecodePluginCreator : public IPluginCreator {
public:
RpnDecodePluginCreator() {}
const char *getPluginName() const override {
const char *getPluginName() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override {
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) TRT_NOEXCEPT override {
return new RpnDecodePlugin(serialData, serialLength);
}
void setPluginNamespace(const char *N) override {}
const PluginFieldCollection *getFieldNames() override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) override { return nullptr; }
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {}
const PluginFieldCollection *getFieldNames() TRT_NOEXCEPT override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) TRT_NOEXCEPT override { return nullptr; }
};
REGISTER_TENSORRT_PLUGIN(RpnDecodePluginCreator);

19
rcnn/RpnNms.cu Normal file → Executable file
View File

@ -1,7 +1,6 @@
#include <cuda.h>
#include <thrust/device_ptr.h>
#include <thrust/gather.h>
#include <thrust/system/cuda/detail/cub/device/device_radix_sort.cuh>
#include <algorithm>
#include <iostream>
@ -12,6 +11,16 @@
#include "RpnNmsPlugin.h"
#include "./cuda_utils.h"
#include "macros.h"
#ifdef CUDA_11
#include <cub/device/device_radix_sort.cuh>
#include <cub/iterator/counting_input_iterator.cuh>
#else
#include <thrust/system/cuda/detail/cub/device/device_radix_sort.cuh>
#include <thrust/system/cuda/detail/cub/iterator/counting_input_iterator.cuh>
namespace cub = thrust::cuda_cub::cub;
#endif
namespace nvinfer1 {
@ -48,7 +57,7 @@ namespace nvinfer1 {
}
int rpnNms(int batch_size,
const void *const *inputs, void **outputs,
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
size_t pre_nms_topk, int post_nms_topk, float nms_thresh,
void *workspace, size_t workspace_size, cudaStream_t stream) {
if (!workspace || !workspace_size) {
@ -59,7 +68,7 @@ namespace nvinfer1 {
workspace_size += get_size_aligned<float>(pre_nms_topk); // scores_sorted
size_t temp_size_sort = 0;
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(
cub::DeviceRadixSort::SortPairsDescending(
static_cast<void*>(nullptr), temp_size_sort,
static_cast<float*>(nullptr),
static_cast<float*>(nullptr),
@ -88,7 +97,7 @@ namespace nvinfer1 {
auto out_boxes = static_cast<float4 *>(outputs[0]) + batch * post_nms_topk;
int num_detections = pre_nms_topk;
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
in_scores, scores_sorted, indices, indices_sorted, num_detections, 0,
sizeof(*scores_sorted) * 8, stream);
@ -100,7 +109,7 @@ namespace nvinfer1 {
indices_sorted, scores_sorted, in_boxes);
// Re-sort with updated scores
thrust::cuda_cub::cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
cub::DeviceRadixSort::SortPairsDescending(workspace, workspace_size,
scores_sorted, scores, indices_sorted, indices, num_detections, 0, sizeof(*scores_sorted) * 8, stream);
// Gather filtered scores, boxes, classes

57
rcnn/RpnNmsPlugin.h Normal file → Executable file
View File

@ -4,6 +4,7 @@
#include <vector>
#include <cassert>
#include "macros.h"
using namespace nvinfer1;
@ -14,7 +15,7 @@ using namespace nvinfer1;
namespace nvinfer1 {
int rpnNms(int batchSize,
const void *const *inputs, void **outputs,
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
size_t pre_nms_topk, int post_nms_topk, float nms_thresh,
void *workspace, size_t workspace_size, cudaStream_t stream);
@ -39,12 +40,12 @@ class RpnNmsPlugin : public IPluginV2Ext {
read(d, _pre_nms_topk);
}
size_t getSerializationSize() const override {
size_t getSerializationSize() const TRT_NOEXCEPT override {
return sizeof(_nms_thresh) + sizeof(_post_nms_topk)
+ sizeof(_pre_nms_topk);
}
void serialize(void *buffer) const override {
void serialize(void *buffer) const TRT_NOEXCEPT override {
char* d = static_cast<char*>(buffer);
write(d, _nms_thresh);
write(d, _post_nms_topk);
@ -69,34 +70,34 @@ class RpnNmsPlugin : public IPluginV2Ext {
this->deserialize(data, length);
}
const char *getPluginType() const override {
const char *getPluginType() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
int getNbOutputs() const override {
int getNbOutputs() const TRT_NOEXCEPT override {
return 1;
}
Dims getOutputDimensions(int index,
const Dims *inputs, int nbInputDims) override {
const Dims *inputs, int nbInputDims) TRT_NOEXCEPT override {
assert(nbInputDims == 2);
assert(index < this->getNbOutputs());
return Dims2(_post_nms_topk, 4);
}
bool supportsFormat(DataType type, PluginFormat format) const override {
bool supportsFormat(DataType type, PluginFormat format) const TRT_NOEXCEPT override {
return type == DataType::kFLOAT && format == PluginFormat::kLINEAR;
}
int initialize() override { return 0; }
int initialize() TRT_NOEXCEPT override { return 0; }
void terminate() override {}
void terminate() TRT_NOEXCEPT override {}
size_t getWorkspaceSize(int maxBatchSize) const override {
size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override {
if (size < 0) {
size = rpnNms(maxBatchSize, nullptr, nullptr, _pre_nms_topk,
_post_nms_topk, _nms_thresh,
@ -106,40 +107,40 @@ class RpnNmsPlugin : public IPluginV2Ext {
}
int enqueue(int batchSize,
const void *const *inputs, void **outputs,
void *workspace, cudaStream_t stream) override {
const void *const *inputs, void *TRT_CONST_ENQUEUE*outputs,
void *workspace, cudaStream_t stream) TRT_NOEXCEPT override {
return rpnNms(batchSize, inputs, outputs, _pre_nms_topk,
_post_nms_topk, _nms_thresh,
workspace, getWorkspaceSize(batchSize), stream);
}
void destroy() override {
void destroy() TRT_NOEXCEPT override {
delete this;
}
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
void setPluginNamespace(const char *N) override {
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {
}
// IPluginV2Ext Methods
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const {
DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override {
assert(index < 1);
return DataType::kFLOAT;
}
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted,
int nbInputs) const {
int nbInputs) const TRT_NOEXCEPT override {
return false;
}
bool canBroadcastInputAcrossBatch(int inputIndex) const { return false; }
bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override { return false; }
void configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast,
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) {
const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) TRT_NOEXCEPT override {
assert(*inputTypes == nvinfer1::DataType::kFLOAT &&
floatFormat == nvinfer1::PluginFormat::kLINEAR);
assert(nbInputs == 2);
@ -147,7 +148,7 @@ class RpnNmsPlugin : public IPluginV2Ext {
_pre_nms_topk = inputDims[0].d[0];
}
IPluginV2Ext *clone() const override {
IPluginV2Ext *clone() const TRT_NOEXCEPT override {
return new RpnNmsPlugin(_nms_thresh, _post_nms_topk, _pre_nms_topk);
}
@ -167,24 +168,24 @@ class RpnNmsPluginCreator : public IPluginCreator {
public:
RpnNmsPluginCreator() {}
const char *getPluginNamespace() const override {
const char *getPluginNamespace() const TRT_NOEXCEPT override {
return PLUGIN_NAMESPACE;
}
const char *getPluginName() const override {
const char *getPluginName() const TRT_NOEXCEPT override {
return PLUGIN_NAME;
}
const char *getPluginVersion() const override {
const char *getPluginVersion() const TRT_NOEXCEPT override {
return PLUGIN_VERSION;
}
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override {
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) TRT_NOEXCEPT override {
return new RpnNmsPlugin(serialData, serialLength);
}
void setPluginNamespace(const char *N) override {}
const PluginFieldCollection *getFieldNames() override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) override { return nullptr; }
void setPluginNamespace(const char *N) TRT_NOEXCEPT override {}
const PluginFieldCollection *getFieldNames() TRT_NOEXCEPT override { return nullptr; }
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) TRT_NOEXCEPT override { return nullptr; }
};
REGISTER_TENSORRT_PLUGIN(RpnNmsPluginCreator);

2
rcnn/backbone.hpp Normal file → Executable file
View File

@ -44,7 +44,7 @@ int group_num = 1) {
auto max_pool2d = network->addPoolingNd(*r1->getOutput(0), PoolingType::kMAX, DimsHW{ 3, 3 });
max_pool2d->setStrideNd(DimsHW{ 2, 2 });
max_pool2d->setPaddingNd(DimsHW{ 1, 1 });
auto mp_dim = max_pool2d->getOutput(0)->getDimensions();
// auto mp_dim = max_pool2d->getOutput(0)->getDimensions();
return max_pool2d;
}

17
rcnn/calibrator.hpp Normal file → Executable file
View File

@ -9,6 +9,7 @@
#include <algorithm>
#include "./cuda_utils.h"
#include "common.hpp"
#include "macros.h"
//! \class Int8EntropyCalibrator2
//!
@ -22,10 +23,10 @@ class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 {
const char* input_blob_name, bool read_cache = true);
virtual ~Int8EntropyCalibrator2();
int getBatchSize() const override;
bool getBatch(void* bindings[], const char* names[], int nbBindings) override;
const void* readCalibrationCache(size_t& length) override;
void writeCalibrationCache(const void* cache, size_t length) override;
int getBatchSize() const TRT_NOEXCEPT override;
bool getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT override;
const void* readCalibrationCache(size_t& length) TRT_NOEXCEPT override;
void writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT override;
private:
int batchsize_;
@ -63,11 +64,11 @@ Int8EntropyCalibrator2::~Int8EntropyCalibrator2() {
CUDA_CHECK(cudaFree(device_input_));
}
int Int8EntropyCalibrator2::getBatchSize() const {
int Int8EntropyCalibrator2::getBatchSize() const TRT_NOEXCEPT {
return batchsize_;
}
bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings) {
bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT {
if (img_idx_ + batchsize_ > static_cast<int>(img_files_.size())) {
return false;
}
@ -97,7 +98,7 @@ bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int
return true;
}
const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) {
const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) TRT_NOEXCEPT {
std::cout << "reading calib cache: " << calib_table_name_ << std::endl;
calib_cache_.clear();
std::ifstream input(calib_table_name_, std::ios::binary);
@ -109,7 +110,7 @@ const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length) {
return length ? calib_cache_.data() : nullptr;
}
void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length) {
void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT {
std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl;
std::ofstream output(calib_table_name_, std::ios::binary);
output.write(reinterpret_cast<const char*>(cache), length);

0
rcnn/common.hpp Normal file → Executable file
View File

0
rcnn/cuda_utils.h Normal file → Executable file
View File

0
rcnn/gen_wts.py Normal file → Executable file
View File

3
rcnn/logging.h Normal file → Executable file
View File

@ -18,6 +18,7 @@
#define TENSORRT_LOGGING_H
#include "NvInferRuntimeCommon.h"
#include "macros.h"
#include <cassert>
#include <ctime>
#include <iomanip>
@ -208,7 +209,7 @@ class Logger : public nvinfer1::ILogger {
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) override {
void log(Severity severity, const char* msg) TRT_NOEXCEPT override {
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}

16
rcnn/macros.h Normal file
View File

@ -0,0 +1,16 @@
#pragma once
#include <NvInfer.h>
#include <cuda.h>
#if CUDA_VERSION >=11000
#define CUDA_11
#endif
#if NV_TENSORRT_MAJOR >= 8
#define TRT_NOEXCEPT noexcept
#define TRT_CONST_ENQUEUE const
#else
#define TRT_NOEXCEPT
#define TRT_CONST_ENQUEUE
#endif

0
rcnn/rcnn.cpp Normal file → Executable file
View File