diff --git a/yolov8/main.cpp b/yolov8/main.cpp index 74820db..83fdac4 100644 --- a/yolov8/main.cpp +++ b/yolov8/main.cpp @@ -83,6 +83,10 @@ void prepare_buffer(ICudaEngine *engine, float **input_buffer_device, float **ou if (cuda_post_process == "c") { *output_buffer_host = new float[kBatchSize * kOutputSize]; } else if (cuda_post_process == "g") { + if (kBatchSize > 1) { + std::cerr << "Do not yet support GPU post processing for multiple batches" << std::endl; + exit(0); + } // Allocate memory for decode_ptr_host and copy to device *decode_ptr_host = new float[1 + kMaxNumOutputBbox * bbox_element]; CUDA_CHECK(cudaMalloc((void **)decode_ptr_device, sizeof(float) * (1 + kMaxNumOutputBbox * bbox_element))); diff --git a/yolov8/plugin/yololayer.cu b/yolov8/plugin/yololayer.cu old mode 100644 new mode 100755 index 1923701..40f1555 --- a/yolov8/plugin/yololayer.cu +++ b/yolov8/plugin/yololayer.cu @@ -2,6 +2,9 @@ #include "types.h" #include #include +#include "cuda_utils.h" +#include +#include namespace Tn { template @@ -18,220 +21,213 @@ namespace Tn { } // namespace Tn -namespace nvinfer1 -{ - YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut) { - mClassCount = classCount; - mYoloV8NetWidth = netWidth; - mYoloV8netHeight = netHeight; - mMaxOutObject = maxOut; - } +namespace nvinfer1 { +YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut) { + mClassCount = classCount; + mYoloV8NetWidth = netWidth; + mYoloV8netHeight = netHeight; + mMaxOutObject = maxOut; +} - YoloLayerPlugin::~YoloLayerPlugin() {} +YoloLayerPlugin::~YoloLayerPlugin() {} - YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length) { - using namespace Tn; - const char* d = reinterpret_cast(data), * a = d; - read(d, mClassCount); - read(d, mThreadCount); - read(d, mYoloV8NetWidth); - read(d, mYoloV8netHeight); - read(d, mMaxOutObject); +YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length) { + using namespace Tn; + const char* d = reinterpret_cast(data), * a = d; + read(d, mClassCount); + read(d, mThreadCount); + read(d, mYoloV8NetWidth); + read(d, mYoloV8netHeight); + read(d, mMaxOutObject); - assert(d == a + length); - } + assert(d == a + length); +} + +void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT { + + using namespace Tn; + char* d = static_cast(buffer), * a = d; + write(d, mClassCount); + write(d, mThreadCount); + write(d, mYoloV8NetWidth); + write(d, mYoloV8netHeight); + write(d, mMaxOutObject); + + assert(d == a + getSerializationSize()); +} + +size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT { + return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mYoloV8netHeight) + sizeof(mYoloV8NetWidth) + sizeof(mMaxOutObject); +} + +int YoloLayerPlugin::initialize() TRT_NOEXCEPT { + return 0; +} + +nvinfer1::Dims YoloLayerPlugin::getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims) TRT_NOEXCEPT { + int total_size = mMaxOutObject * sizeof(Detection) / sizeof(float); + return nvinfer1::Dims3(total_size + 1, 1, 1); +} + +void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT { + mPluginNamespace = pluginNamespace; +} + +const char* YoloLayerPlugin::getPluginNamespace() const TRT_NOEXCEPT { + return mPluginNamespace; +} + +nvinfer1::DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT { + return nvinfer1::DataType::kFLOAT; +} + +bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT { + + return false; +} + +bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT { + + return false; +} + +void YoloLayerPlugin::configurePlugin(nvinfer1::PluginTensorDesc const* in, int nbInput, nvinfer1::PluginTensorDesc const* out, int nbOutput) TRT_NOEXCEPT {}; + +void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT {}; + +void YoloLayerPlugin::detachFromContext() TRT_NOEXCEPT {} + +const char* YoloLayerPlugin::getPluginType() const TRT_NOEXCEPT { + + return "YoloLayer_TRT"; +} + +const char* YoloLayerPlugin::getPluginVersion() const TRT_NOEXCEPT { + return "1"; +} + +void YoloLayerPlugin::destroy() TRT_NOEXCEPT { + + delete this; +} + +nvinfer1::IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT { + + YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV8NetWidth, mYoloV8netHeight, mMaxOutObject); + p->setPluginNamespace(mPluginNamespace); + return p; +} + +int YoloLayerPlugin::enqueue(int batchSize, const void* TRT_CONST_ENQUEUE* inputs, void* const* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT { + + forwardGpu((const float* const*)inputs, (float*)outputs[0], stream, mYoloV8netHeight, mYoloV8NetWidth, batchSize); + return 0; +} - void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT { - - using namespace Tn; - char* d = static_cast(buffer), * a = d; - write(d, mClassCount); - write(d, mThreadCount); - write(d, mYoloV8NetWidth); - write(d, mYoloV8netHeight); - write(d, mMaxOutObject); +__device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); }; - assert(d == a + getSerializationSize()); - } +__global__ void CalDetection(const float* input, float* output, int numElements, int maxoutobject, + const int grid_h, int grid_w, const int stride, int classes, int outputElem) { + int idx = threadIdx.x + blockDim.x * blockIdx.x; + if (idx >= numElements) return; - size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT { - return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mYoloV8netHeight) + sizeof(mYoloV8NetWidth) + sizeof(mMaxOutObject); - } + int total_grid = grid_h * grid_w; + int info_len = 4 + classes; + int batchIdx = idx / total_grid; + int elemIdx = idx % total_grid; + const float* curInput = input + batchIdx * total_grid * info_len; + int outputIdx = batchIdx * outputElem; - int YoloLayerPlugin::initialize() TRT_NOEXCEPT { - return 0; - } - - nvinfer1::Dims YoloLayerPlugin::getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims) TRT_NOEXCEPT { - int total_size = mMaxOutObject * sizeof(Detection) / sizeof(float); - return nvinfer1::Dims3(total_size + 1, 1, 1); - } - - void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT { - mPluginNamespace = pluginNamespace; - } - - const char* YoloLayerPlugin::getPluginNamespace() const TRT_NOEXCEPT { - return mPluginNamespace; - } - - nvinfer1::DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT { - return nvinfer1::DataType::kFLOAT; - } - - - bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT { - - return false; - } - - bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT { - - return false; - } - - - void YoloLayerPlugin::configurePlugin(nvinfer1::PluginTensorDesc const* in, int nbInput, nvinfer1::PluginTensorDesc const* out, int nbOutput) TRT_NOEXCEPT {}; - - void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT {}; - - void YoloLayerPlugin::detachFromContext() TRT_NOEXCEPT {} - - const char* YoloLayerPlugin::getPluginType() const TRT_NOEXCEPT { - - return "YoloLayer_TRT"; - } - - const char* YoloLayerPlugin::getPluginVersion() const TRT_NOEXCEPT { - return "1"; - } - - void YoloLayerPlugin::destroy() TRT_NOEXCEPT { - - delete this; - } - - nvinfer1::IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT - - { - - YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV8NetWidth, mYoloV8netHeight, mMaxOutObject); - p->setPluginNamespace(mPluginNamespace); - return p; - } - - - int YoloLayerPlugin::enqueue(int batchSize, const void* TRT_CONST_ENQUEUE* inputs, void* const* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT { - - forwardGpu((const float* const*)inputs, (float*)outputs[0], stream, mYoloV8netHeight, mYoloV8NetWidth, batchSize); - return 0; - } - - - __device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); }; - - - __global__ void CalDetection(const float* input, float* output, int numElements, int maxoutobject, const int grid_h, int grid_w, const int stride, int classes) { - int idx = threadIdx.x + blockDim.x * blockIdx.x; - if (idx >= numElements) return; - - int total_grid = grid_h * grid_w; - int info_len = 4 + classes; - const float* curInput = input; - - int class_id = 0; - float max_cls_prob = 0.0; - for (int i = 4; i < info_len; i++) { - float p = Logist(curInput[idx + i * total_grid]); - if (p > max_cls_prob) { - max_cls_prob = p; - class_id = i - 4; - } - } - - if (max_cls_prob < 0.1) return; - - int count = (int)atomicAdd(output, 1); - if (count >= maxoutobject) return; - char* data = (char*)output + sizeof(float) + count * sizeof(Detection); - Detection* det = (Detection*)(data); - - int row = idx / grid_w; - int col = idx % grid_w; - - det->conf = max_cls_prob; - det->class_id = class_id; - det->bbox[0] = (col + 0.5f - curInput[idx + 0 * total_grid]) * stride; - det->bbox[1] = (row + 0.5f - curInput[idx + 1 * total_grid]) * stride; - det->bbox[2] = (col + 0.5f + curInput[idx + 2 * total_grid]) * stride; - det->bbox[3] = (row + 0.5f + curInput[idx + 3 * total_grid]) * stride; - } - - - - void YoloLayerPlugin::forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int mYoloV8netHeight,int mYoloV8NetWidth, int batchSize) { - int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float); - cudaMemsetAsync(output, 0, sizeof(float), stream); - - int numElem = 0; - int grids[3][2] = { {mYoloV8netHeight / 8, mYoloV8NetWidth / 8}, {mYoloV8netHeight / 16, mYoloV8NetWidth / 16}, {mYoloV8netHeight / 32, mYoloV8NetWidth / 32} }; - int strides[] = { 8, 16, 32 }; - for (unsigned int i = 0; i < 3; i++) { - int grid_h = grids[i][0]; - int grid_w = grids[i][1]; - int stride = strides[i]; - numElem = grid_h * grid_w; - if (numElem < mThreadCount) mThreadCount = numElem; - - CalDetection << <(numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> > - (inputs[i], output, numElem, mMaxOutObject, grid_h, grid_w, stride, mClassCount); + int class_id = 0; + float max_cls_prob = 0.0; + for (int i = 4; i < info_len; i++) { + float p = Logist(curInput[elemIdx + i * total_grid]); + if (p > max_cls_prob) { + max_cls_prob = p; + class_id = i - 4; } } - - PluginFieldCollection YoloPluginCreator::mFC{}; - std::vector YoloPluginCreator::mPluginAttributes; + if (max_cls_prob < 0.1) return; - YoloPluginCreator::YoloPluginCreator() { - mPluginAttributes.clear(); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); + int count = (int)atomicAdd(output + outputIdx, 1); + if (count >= maxoutobject) return; + char* data = (char*)(output + outputIdx) + sizeof(float) + count * sizeof(Detection); + Detection* det = (Detection*)(data); + + int row = elemIdx / grid_w; + int col = elemIdx % grid_w; + + det->conf = max_cls_prob; + det->class_id = class_id; + det->bbox[0] = (col + 0.5f - curInput[elemIdx + 0 * total_grid]) * stride; + det->bbox[1] = (row + 0.5f - curInput[elemIdx + 1 * total_grid]) * stride; + det->bbox[2] = (col + 0.5f + curInput[elemIdx + 2 * total_grid]) * stride; + det->bbox[3] = (row + 0.5f + curInput[elemIdx + 3 * total_grid]) * stride; +} + +void YoloLayerPlugin::forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int mYoloV8netHeight,int mYoloV8NetWidth, int batchSize) { + int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float); + cudaMemsetAsync(output, 0, sizeof(float), stream); + for (int idx = 0; idx < batchSize; ++idx) { + CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream)); } + int numElem = 0; + int grids[3][2] = { {mYoloV8netHeight / 8, mYoloV8NetWidth / 8}, {mYoloV8netHeight / 16, mYoloV8NetWidth / 16}, {mYoloV8netHeight / 32, mYoloV8NetWidth / 32} }; + int strides[] = { 8, 16, 32 }; + for (unsigned int i = 0; i < 3; i++) { + int grid_h = grids[i][0]; + int grid_w = grids[i][1]; + int stride = strides[i]; + numElem = grid_h * grid_w * batchSize; + if (numElem < mThreadCount) mThreadCount = numElem; - const char* YoloPluginCreator::getPluginName() const TRT_NOEXCEPT { - return "YoloLayer_TRT"; + CalDetection << <(numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> > + (inputs[i], output, numElem, mMaxOutObject, grid_h, grid_w, stride, mClassCount, outputElem); } +} - const char* YoloPluginCreator::getPluginVersion() const TRT_NOEXCEPT { - return "1"; - } +PluginFieldCollection YoloPluginCreator::mFC{}; +std::vector YoloPluginCreator::mPluginAttributes; - const PluginFieldCollection* YoloPluginCreator::getFieldNames() TRT_NOEXCEPT { - return &mFC; - } +YoloPluginCreator::YoloPluginCreator() { + mPluginAttributes.clear(); + mFC.nbFields = mPluginAttributes.size(); + mFC.fields = mPluginAttributes.data(); +} - IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT { - assert(fc->nbFields == 1); - assert(strcmp(fc->fields[0].name, "netinfo") == 0); - int* p_netinfo = (int*)(fc->fields[0].data); - int class_count = p_netinfo[0]; - int input_w = p_netinfo[1]; - int input_h = p_netinfo[2]; - int max_output_object_count = p_netinfo[3]; +const char* YoloPluginCreator::getPluginName() const TRT_NOEXCEPT { + return "YoloLayer_TRT"; +} - YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } +const char* YoloPluginCreator::getPluginVersion() const TRT_NOEXCEPT { + return "1"; +} +const PluginFieldCollection* YoloPluginCreator::getFieldNames() TRT_NOEXCEPT { + return &mFC; +} - IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT { - // This object will be deleted when the network is destroyed, which will - // call YoloLayerPlugin::destroy() - YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } +IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT { + assert(fc->nbFields == 1); + assert(strcmp(fc->fields[0].name, "netinfo") == 0); + int* p_netinfo = (int*)(fc->fields[0].data); + int class_count = p_netinfo[0]; + int input_w = p_netinfo[1]; + int input_h = p_netinfo[2]; + int max_output_object_count = p_netinfo[3]; + YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; +} + +IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT { + // This object will be deleted when the network is destroyed, which will + // call YoloLayerPlugin::destroy() + YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; +} } // namespace nvinfer1 diff --git a/yolov8/src/model.cpp b/yolov8/src/model.cpp index 54b2063..37db0d4 100644 --- a/yolov8/src/model.cpp +++ b/yolov8/src/model.cpp @@ -72,8 +72,8 @@ nvinfer1::IHostMemory* buildEngineYolov8n(nvinfer1::IBuilder* builder, conv22_cv2_0_2->setStrideNd(nvinfer1::DimsHW{1, 1}); conv22_cv2_0_2->setPaddingNd(nvinfer1::DimsHW{0, 0}); - nvinfer1::IElementWiseLayer* conv22_cv3_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 64, 3, 1, 1, "model.22.cv3.0.0"); - nvinfer1::IElementWiseLayer* conv22_cv3_0_1 = convBnSiLU(network, weightMap, *conv22_cv3_0_0->getOutput(0), 64, 3, 1, 1, "model.22.cv3.0.1"); + nvinfer1::IElementWiseLayer* conv22_cv3_0_0 = convBnSiLU(network, weightMap, *conv15->getOutput(0), 80, 3, 1, 1, "model.22.cv3.0.0"); + nvinfer1::IElementWiseLayer* conv22_cv3_0_1 = convBnSiLU(network, weightMap, *conv22_cv3_0_0->getOutput(0), 80, 3, 1, 1, "model.22.cv3.0.1"); nvinfer1::IConvolutionLayer* conv22_cv3_0_2 = network->addConvolutionNd(*conv22_cv3_0_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1,1}, weightMap["model.22.cv3.0.2.weight"], weightMap["model.22.cv3.0.2.bias"]); conv22_cv3_0_2->setStride(nvinfer1::DimsHW{1, 1}); conv22_cv3_0_2->setPadding(nvinfer1::DimsHW{0, 0}); @@ -87,8 +87,8 @@ nvinfer1::IHostMemory* buildEngineYolov8n(nvinfer1::IBuilder* builder, conv22_cv2_1_2->setStrideNd(nvinfer1::DimsHW{1,1}); conv22_cv2_1_2->setPaddingNd(nvinfer1::DimsHW{0,0}); - nvinfer1::IElementWiseLayer* conv22_cv3_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 64, 3, 1, 1, "model.22.cv3.1.0"); - nvinfer1::IElementWiseLayer* conv22_cv3_1_1 = convBnSiLU(network, weightMap, *conv22_cv3_1_0->getOutput(0), 64, 3, 1, 1, "model.22.cv3.1.1"); + nvinfer1::IElementWiseLayer* conv22_cv3_1_0 = convBnSiLU(network, weightMap, *conv18->getOutput(0), 80, 3, 1, 1, "model.22.cv3.1.0"); + nvinfer1::IElementWiseLayer* conv22_cv3_1_1 = convBnSiLU(network, weightMap, *conv22_cv3_1_0->getOutput(0), 80, 3, 1, 1, "model.22.cv3.1.1"); nvinfer1::IConvolutionLayer* conv22_cv3_1_2 = network->addConvolutionNd(*conv22_cv3_1_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1, 1}, weightMap["model.22.cv3.1.2.weight"], weightMap["model.22.cv3.1.2.bias"]); conv22_cv3_1_2->setStrideNd(nvinfer1::DimsHW{1,1}); conv22_cv3_1_2->setPaddingNd(nvinfer1::DimsHW{0,0}); @@ -101,8 +101,8 @@ nvinfer1::IHostMemory* buildEngineYolov8n(nvinfer1::IBuilder* builder, nvinfer1::IElementWiseLayer* conv22_cv2_2_1 = convBnSiLU(network, weightMap, *conv22_cv2_2_0->getOutput(0), 64, 3, 1, 1, "model.22.cv2.2.1"); nvinfer1::IConvolutionLayer* conv22_cv2_2_2 = network->addConvolution(*conv22_cv2_2_1->getOutput(0), 64, nvinfer1::DimsHW{1,1}, weightMap["model.22.cv2.2.2.weight"], weightMap["model.22.cv2.2.2.bias"]); - nvinfer1::IElementWiseLayer* conv22_cv3_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 64, 3, 1, 1, "model.22.cv3.2.0"); - nvinfer1::IElementWiseLayer* conv22_cv3_2_1 = convBnSiLU(network, weightMap, *conv22_cv3_2_0->getOutput(0), 64, 3, 1, 1, "model.22.cv3.2.1"); + nvinfer1::IElementWiseLayer* conv22_cv3_2_0 = convBnSiLU(network, weightMap, *conv21->getOutput(0), 80, 3, 1, 1, "model.22.cv3.2.0"); + nvinfer1::IElementWiseLayer* conv22_cv3_2_1 = convBnSiLU(network, weightMap, *conv22_cv3_2_0->getOutput(0), 80, 3, 1, 1, "model.22.cv3.2.1"); nvinfer1::IConvolutionLayer* conv22_cv3_2_2 = network->addConvolution(*conv22_cv3_2_1->getOutput(0), kNumClass, nvinfer1::DimsHW{1,1}, weightMap["model.22.cv3.2.2.weight"], weightMap["model.22.cv3.2.2.bias"]); nvinfer1::ITensor* inputTensor22_2[] = {conv22_cv2_2_2->getOutput(0), conv22_cv3_2_2->getOutput(0)}; @@ -795,4 +795,4 @@ nvinfer1::IHostMemory* buildEngineYolov8x(nvinfer1::IBuilder* builder, free((void*)(mem.second.values)); } return serialized_model; -} \ No newline at end of file +}