From 8bdbb47dfafc4943bbba4ee9ee160d4f89ba457d Mon Sep 17 00:00:00 2001 From: BaofengZan Date: Fri, 25 Sep 2020 10:11:57 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BB=8Ewts=E6=96=87=E4=BB=B6=E4=B8=AD?= =?UTF-8?q?=E8=AF=BB=E5=8F=96anhor=E4=BF=A1=E6=81=AF=EF=BC=8C=E5=B9=B6?= =?UTF-8?q?=E4=BC=A0=E5=85=A5yololayer=20plugin=20(#215)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 修改图片地址为URL * get Anchor values from wts * get anchor values form wts * get anchor values from wts * get anchor values from wts * 1 using space 2 release heap mem * 1 Using 4 space 2 Delete Chinese comment 3 Load all variables 4 Release heap mem * 1 Using 4 space 2 Delete Chinese comment 3 Load all variables 4 Release heap mem * 1. 函数模块化 2. 格式化代码 3. 同步为最新的yolov5代码 --- yolov5/common.hpp | 168 ++++++++++++++++++++++++++++------------ yolov5/yololayer.cu | 185 ++++++++++++++++++++++++++++++-------------- yolov5/yololayer.h | 156 +++++++++++++++++-------------------- yolov5/yolov5.cpp | 109 ++++++++------------------ 4 files changed, 351 insertions(+), 267 deletions(-) diff --git a/yolov5/common.hpp b/yolov5/common.hpp index d31f7fc..c42c8cf 100644 --- a/yolov5/common.hpp +++ b/yolov5/common.hpp @@ -32,8 +32,9 @@ cv::Mat preprocess_img(cv::Mat& img) { h = r_w * img.rows; x = 0; y = (Yolo::INPUT_H - h) / 2; - } else { - w = r_h* img.cols; + } + else { + w = r_h * img.cols; h = Yolo::INPUT_H; x = (Yolo::INPUT_W - w) / 2; y = 0; @@ -50,40 +51,41 @@ cv::Rect get_rect(cv::Mat& img, float bbox[4]) { float r_w = Yolo::INPUT_W / (img.cols * 1.0); float r_h = Yolo::INPUT_H / (img.rows * 1.0); if (r_h > r_w) { - l = bbox[0] - bbox[2]/2.f; - r = bbox[0] + bbox[2]/2.f; - t = bbox[1] - bbox[3]/2.f - (Yolo::INPUT_H - r_w * img.rows) / 2; - b = bbox[1] + bbox[3]/2.f - (Yolo::INPUT_H - r_w * img.rows) / 2; + l = bbox[0] - bbox[2] / 2.f; + r = bbox[0] + bbox[2] / 2.f; + t = bbox[1] - bbox[3] / 2.f - (Yolo::INPUT_H - r_w * img.rows) / 2; + b = bbox[1] + bbox[3] / 2.f - (Yolo::INPUT_H - r_w * img.rows) / 2; l = l / r_w; r = r / r_w; t = t / r_w; b = b / r_w; - } else { - l = bbox[0] - bbox[2]/2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; - r = bbox[0] + bbox[2]/2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; - t = bbox[1] - bbox[3]/2.f; - b = bbox[1] + bbox[3]/2.f; + } + else { + l = bbox[0] - bbox[2] / 2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; + r = bbox[0] + bbox[2] / 2.f - (Yolo::INPUT_W - r_h * img.cols) / 2; + t = bbox[1] - bbox[3] / 2.f; + b = bbox[1] + bbox[3] / 2.f; l = l / r_h; r = r / r_h; t = t / r_h; b = b / r_h; } - return cv::Rect(l, t, r-l, b-t); + return cv::Rect(l, t, r - l, b - t); } float iou(float lbox[4], float rbox[4]) { float interBox[] = { - std::max(lbox[0] - lbox[2]/2.f , rbox[0] - rbox[2]/2.f), //left - std::min(lbox[0] + lbox[2]/2.f , rbox[0] + rbox[2]/2.f), //right - std::max(lbox[1] - lbox[3]/2.f , rbox[1] - rbox[3]/2.f), //top - std::min(lbox[1] + lbox[3]/2.f , rbox[1] + rbox[3]/2.f), //bottom + (std::max)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left + (std::min)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right + (std::max)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top + (std::min)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom }; - if(interBox[2] > interBox[3] || interBox[0] > interBox[1]) + if (interBox[2] > interBox[3] || interBox[0] > interBox[1]) return 0.0f; - float interBoxS =(interBox[1]-interBox[0])*(interBox[3]-interBox[2]); - return interBoxS/(lbox[2]*lbox[3] + rbox[2]*rbox[3] -interBoxS); + float interBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]); + return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS); } bool cmp(const Yolo::Detection& a, const Yolo::Detection& b) { @@ -109,7 +111,7 @@ void nms(std::vector& res, float *output, float conf_thresh, fl res.push_back(item); for (size_t n = m + 1; n < dets.size(); ++n) { if (iou(item.bbox, dets[n].bbox) > nms_thresh) { - dets.erase(dets.begin()+n); + dets.erase(dets.begin() + n); --n; } } @@ -134,7 +136,7 @@ std::map loadWeights(const std::string file) { while (count--) { - Weights wt{DataType::kFLOAT, nullptr, 0}; + Weights wt{ DataType::kFLOAT, nullptr, 0 }; uint32_t size; // Read name and type of blob @@ -149,7 +151,7 @@ std::map loadWeights(const std::string file) { input >> std::hex >> val[x]; } wt.values = val; - + wt.count = size; weightMap[name] = wt; } @@ -168,19 +170,19 @@ IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map(malloc(sizeof(float) * len)); for (int i = 0; i < len; i++) { shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); } - Weights shift{DataType::kFLOAT, shval, len}; + Weights shift{ DataType::kFLOAT, shval, len }; float *pval = reinterpret_cast(malloc(sizeof(float) * len)); for (int i = 0; i < len; i++) { pval[i] = 1.0; } - Weights power{DataType::kFLOAT, pval, len}; + Weights power{ DataType::kFLOAT, pval, len }; weightMap[lname + ".scale"] = scale; weightMap[lname + ".shift"] = shift; @@ -191,12 +193,12 @@ IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname) { - Weights emptywts{DataType::kFLOAT, nullptr, 0}; + Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; int p = ksize / 2; - IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[lname + ".conv.weight"], emptywts); + IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts); assert(conv1); - conv1->setStrideNd(DimsHW{s, s}); - conv1->setPaddingNd(DimsHW{p, p}); + conv1->setStrideNd(DimsHW{ s, s }); + conv1->setPaddingNd(DimsHW{ p, p }); conv1->setNbGroups(g); IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-3); @@ -211,11 +213,11 @@ ILayer* convBlock(INetworkDefinition *network, std::map& w } ILayer* focus(INetworkDefinition *network, std::map& weightMap, ITensor& input, int inch, int outch, int ksize, std::string lname) { - ISliceLayer *s1 = network->addSlice(input, Dims3{0, 0, 0}, Dims3{inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2}, Dims3{1, 2, 2}); - ISliceLayer *s2 = network->addSlice(input, Dims3{0, 1, 0}, Dims3{inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2}, Dims3{1, 2, 2}); - ISliceLayer *s3 = network->addSlice(input, Dims3{0, 0, 1}, Dims3{inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2}, Dims3{1, 2, 2}); - ISliceLayer *s4 = network->addSlice(input, Dims3{0, 1, 1}, Dims3{inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2}, Dims3{1, 2, 2}); - ITensor* inputTensors[] = {s1->getOutput(0), s2->getOutput(0), s3->getOutput(0), s4->getOutput(0)}; + ISliceLayer *s1 = network->addSlice(input, Dims3{ 0, 0, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); + ISliceLayer *s2 = network->addSlice(input, Dims3{ 0, 1, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); + ISliceLayer *s3 = network->addSlice(input, Dims3{ 0, 0, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); + ISliceLayer *s4 = network->addSlice(input, Dims3{ 0, 1, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); + ITensor* inputTensors[] = { s1->getOutput(0), s2->getOutput(0), s3->getOutput(0), s4->getOutput(0) }; auto cat = network->addConcatenation(inputTensors, 4); auto conv = convBlock(network, weightMap, *cat->getOutput(0), outch, ksize, 1, 1, lname + ".conv"); return conv; @@ -232,18 +234,18 @@ ILayer* bottleneck(INetworkDefinition *network, std::map& } ILayer* bottleneckCSP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { - Weights emptywts{DataType::kFLOAT, nullptr, 0}; + Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; int c_ = (int)((float)c2 * e); auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); - auto cv2 = network->addConvolutionNd(input, c_, DimsHW{1, 1}, weightMap[lname + ".cv2.weight"], emptywts); + auto cv2 = network->addConvolutionNd(input, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv2.weight"], emptywts); ITensor *y1 = cv1->getOutput(0); for (int i = 0; i < n; i++) { auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i)); y1 = b->getOutput(0); } - auto cv3 = network->addConvolutionNd(*y1, c_, DimsHW{1, 1}, weightMap[lname + ".cv3.weight"], emptywts); + auto cv3 = network->addConvolutionNd(*y1, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv3.weight"], emptywts); - ITensor* inputTensors[] = {cv3->getOutput(0), cv2->getOutput(0)}; + ITensor* inputTensors[] = { cv3->getOutput(0), cv2->getOutput(0) }; auto cat = network->addConcatenation(inputTensors, 2); IScaleLayer* bn = addBatchNorm2d(network, weightMap, *cat->getOutput(0), lname + ".bn", 1e-4); @@ -258,17 +260,17 @@ ILayer* SPP(INetworkDefinition *network, std::map& weightM int c_ = c1 / 2; auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); - auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{k1, k1}); - pool1->setPaddingNd(DimsHW{k1 / 2, k1 / 2}); - pool1->setStrideNd(DimsHW{1, 1}); - auto pool2 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{k2, k2}); - pool2->setPaddingNd(DimsHW{k2 / 2, k2 / 2}); - pool2->setStrideNd(DimsHW{1, 1}); - auto pool3 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{k3, k3}); - pool3->setPaddingNd(DimsHW{k3 / 2, k3 / 2}); - pool3->setStrideNd(DimsHW{1, 1}); + auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k1, k1 }); + pool1->setPaddingNd(DimsHW{ k1 / 2, k1 / 2 }); + pool1->setStrideNd(DimsHW{ 1, 1 }); + auto pool2 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k2, k2 }); + pool2->setPaddingNd(DimsHW{ k2 / 2, k2 / 2 }); + pool2->setStrideNd(DimsHW{ 1, 1 }); + auto pool3 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k3, k3 }); + pool3->setPaddingNd(DimsHW{ k3 / 2, k3 / 2 }); + pool3->setStrideNd(DimsHW{ 1, 1 }); - ITensor* inputTensors[] = {cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0)}; + ITensor* inputTensors[] = { cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0) }; auto cat = network->addConcatenation(inputTensors, 4); auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2"); @@ -284,7 +286,7 @@ int read_files_in_dir(const char *p_dir_name, std::vector &file_nam struct dirent* p_file = nullptr; while ((p_file = readdir(p_dir)) != nullptr) { if (strcmp(p_file->d_name, ".") != 0 && - strcmp(p_file->d_name, "..") != 0) { + strcmp(p_file->d_name, "..") != 0) { //std::string cur_file_name(p_dir_name); //cur_file_name += "/"; //cur_file_name += p_file->d_name; @@ -297,5 +299,71 @@ int read_files_in_dir(const char *p_dir_name, std::vector &file_nam return 0; } +std::vector GetAnchors(std::map& weightMap) +{ + std::vector anchors_yolo; + Weights Yolo_Anchors = weightMap["model.24.anchor_grid"]; + assert(Yolo_Anchors.count == 18); + int each_yololayer_anchorsnum = Yolo_Anchors.count / 3; + const float* tempAnchors = (const float*)(Yolo_Anchors.values); + for (int i = 0; i < Yolo_Anchors.count; i++) + { + if (i < each_yololayer_anchorsnum) + { + anchors_yolo.push_back(const_cast(tempAnchors)[i]); + } + if ((i >= each_yololayer_anchorsnum) && (i < (2 * each_yololayer_anchorsnum))) + { + anchors_yolo.push_back(const_cast(tempAnchors)[i]); + } + if (i >= (2 * each_yololayer_anchorsnum)) + { + anchors_yolo.push_back(const_cast(tempAnchors)[i]); + } + } + return anchors_yolo; +} + +IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map& weightMap, IConvolutionLayer* det0, IConvolutionLayer* det1, IConvolutionLayer* det2) +{ + auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); + std::vector anchors_yolo = GetAnchors(weightMap); + PluginField pluginMultidata[4]; + int* NetData = new int[4]; + NetData[0] = Yolo::CLASS_NUM; + NetData[1] = Yolo::INPUT_W; + NetData[2] = Yolo::INPUT_H; + NetData[3] = Yolo::MAX_OUTPUT_BBOX_COUNT; + pluginMultidata[0].data = NetData; + pluginMultidata[0].length = 3; + std::string name = "netdata"; + pluginMultidata[0].name = new char[name.size() + 1]; + strcpy(const_cast(pluginMultidata[0].name), name.c_str()); + pluginMultidata[0].type = PluginFieldType::kFLOAT32; + int scale[3] = { 8, 16, 32 }; + for (int k = 1; k < 4; k++) + { + int* plugindata = new int[8]; + plugindata[0] = Yolo::INPUT_W / scale[k - 1]; + plugindata[1] = Yolo::INPUT_H / scale[k - 1]; + for (int i = 2; i < 8; i++) + { + plugindata[i] = int(anchors_yolo[(k - 1) * 6 + i - 2]); + } + pluginMultidata[k].data = plugindata; + pluginMultidata[k].length = 8; + std::string name = "yolodata" + std::to_string(k); + pluginMultidata[k].name = new char[name.size() + 1]; + strcpy(const_cast(pluginMultidata[k].name), name.c_str()); + pluginMultidata[k].type = PluginFieldType::kFLOAT32; + } + PluginFieldCollection pluginData; + pluginData.nbFields = 4; + pluginData.fields = pluginMultidata; + IPluginV2 *pluginObj = creator->createPlugin("yololayer", &pluginData); + ITensor* inputTensors_yolo[] = { det2->getOutput(0), det1->getOutput(0), det0->getOutput(0) }; + auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj); + return yolo; +} #endif diff --git a/yolov5/yololayer.cu b/yolov5/yololayer.cu index e14ee18..f41e763 100644 --- a/yolov5/yololayer.cu +++ b/yolov5/yololayer.cu @@ -6,29 +6,27 @@ using namespace Yolo; namespace nvinfer1 { - YoloLayerPlugin::YoloLayerPlugin() + YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector& vYoloKernel) { - mClassCount = CLASS_NUM; - mYoloKernel.clear(); - mYoloKernel.push_back(yolo1); - mYoloKernel.push_back(yolo2); - mYoloKernel.push_back(yolo3); - - mKernelCount = mYoloKernel.size(); + mClassCount = classCount; + mYoloV5NetWidth = netWidth; + mYoloV5NetHeight = netHeight; + mMaxOutObject = maxOut; + mYoloKernel = vYoloKernel; + mKernelCount = vYoloKernel.size(); CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); - size_t AnchorLen = sizeof(float)* CHECK_COUNT*2; - for(int ii = 0; ii < mKernelCount; ii ++) + size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2; + for (int ii = 0; ii < mKernelCount; ii++) { - CUDA_CHECK(cudaMalloc(&mAnchor[ii],AnchorLen)); + CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); const auto& yolo = mYoloKernel[ii]; CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); } } - YoloLayerPlugin::~YoloLayerPlugin() { - for(int ii = 0; ii < mKernelCount; ii ++) + for (int ii = 0; ii < mKernelCount; ii++) { CUDA_CHECK(cudaFree(mAnchor[ii])); } @@ -43,20 +41,21 @@ namespace nvinfer1 read(d, mClassCount); read(d, mThreadCount); read(d, mKernelCount); + read(d, mYoloV5NetWidth); + read(d, mYoloV5NetHeight); + read(d, mMaxOutObject); mYoloKernel.resize(mKernelCount); - auto kernelSize = mKernelCount*sizeof(YoloKernel); - memcpy(mYoloKernel.data(),d,kernelSize); + auto kernelSize = mKernelCount * sizeof(YoloKernel); + memcpy(mYoloKernel.data(), d, kernelSize); d += kernelSize; - CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); - size_t AnchorLen = sizeof(float)* CHECK_COUNT*2; - for(int ii = 0; ii < mKernelCount; ii ++) + size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2; + for (int ii = 0; ii < mKernelCount; ii++) { - CUDA_CHECK(cudaMalloc(&mAnchor[ii],AnchorLen)); + CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); const auto& yolo = mYoloKernel[ii]; CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); } - assert(d == a + length); } @@ -67,27 +66,30 @@ namespace nvinfer1 write(d, mClassCount); write(d, mThreadCount); write(d, mKernelCount); - auto kernelSize = mKernelCount*sizeof(YoloKernel); - memcpy(d,mYoloKernel.data(),kernelSize); + write(d, mYoloV5NetWidth); + write(d, mYoloV5NetHeight); + write(d, mMaxOutObject); + auto kernelSize = mKernelCount * sizeof(YoloKernel); + memcpy(d, mYoloKernel.data(), kernelSize); d += kernelSize; assert(d == a + getSerializationSize()); } - + size_t YoloLayerPlugin::getSerializationSize() const - { - return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size(); + { + return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV5NetWidth) + sizeof(mYoloV5NetHeight) + sizeof(mMaxOutObject); } int YoloLayerPlugin::initialize() - { + { return 0; } - + Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) { //output the result to channel - int totalsize = MAX_OUTPUT_BBOX_COUNT * sizeof(Detection) / sizeof(float); + int totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float); return Dims3(totalsize + 1, 1, 1); } @@ -151,25 +153,29 @@ namespace nvinfer1 // Clone the plugin IPluginV2IOExt* YoloLayerPlugin::clone() const { - YoloLayerPlugin *p = new YoloLayerPlugin(); + //YoloLayerPlugin *p = nullptr; + //p = new YoloLayerPlugin(); + YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, mYoloKernel); p->setPluginNamespace(mPluginNamespace); return p; } - __device__ float Logist(float data){ return 1.0f / (1.0f + expf(-data)); }; + __device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); }; + + __global__ void CalDetection(const float *input, float *output, int noElements, + const int netwidth, const int netheight, int maxoutobject, int yoloWidth, int yoloHeight, const float anchors[CHECK_COUNT * 2], int classes, int outputElem) + { - __global__ void CalDetection(const float *input, float *output,int noElements, - int yoloWidth,int yoloHeight,const float anchors[CHECK_COUNT*2],int classes,int outputElem) { - int idx = threadIdx.x + blockDim.x * blockIdx.x; if (idx >= noElements) return; int total_grid = yoloWidth * yoloHeight; int bnIdx = idx / total_grid; - idx = idx - total_grid*bnIdx; + idx = idx - total_grid * bnIdx; int info_len_i = 5 + classes; const float* curInput = input + bnIdx * (info_len_i * total_grid * CHECK_COUNT); + for (int k = 0; k < 3; ++k) { float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]); if (box_prob < IGNORE_THRESH) continue; @@ -182,45 +188,53 @@ namespace nvinfer1 class_id = i - 5; } } - float *res_count = output + bnIdx*outputElem; + float *res_count = output + bnIdx * outputElem; int count = (int)atomicAdd(res_count, 1); - if (count >= MAX_OUTPUT_BBOX_COUNT) return; + if (count >= maxoutobject) return; char* data = (char *)res_count + sizeof(float) + count * sizeof(Detection); - Detection* det = (Detection*)(data); + Detection* det = (Detection*)(data); int row = idx / yoloWidth; int col = idx % yoloWidth; //Location - det->bbox[0] = (col - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * INPUT_W / yoloWidth; - det->bbox[1] = (row - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * INPUT_H / yoloHeight; + // pytorch： + // y = x[i].sigmoid() + // y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy + // y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh + //X: (sigmoid(tx) + cx)/FeaturemapW * netwidth + det->bbox[0] = (col - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth; + det->bbox[1] = (row - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight; + + // W: (Pw * e^tw) / FeaturemapW * netwidth + // v5：https://github.com/ultralytics/yolov5/issues/471 det->bbox[2] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]); - det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2*k]; + det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2 * k]; // det->bbox[3] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]); - det->bbox[3] = det->bbox[3] * det->bbox[3] * anchors[2*k + 1]; + det->bbox[3] = det->bbox[3] * det->bbox[3] * anchors[2 * k + 1]; det->conf = box_prob * max_cls_prob; det->class_id = class_id; } } - void YoloLayerPlugin::forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize) { - - int outputElem = 1 + MAX_OUTPUT_BBOX_COUNT * sizeof(Detection) / sizeof(float); - - for(int idx = 0 ; idx < batchSize; ++idx) { - CUDA_CHECK(cudaMemset(output + idx*outputElem, 0, sizeof(float))); + void YoloLayerPlugin::forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize) + { + int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float); + for (int idx = 0; idx < batchSize; ++idx) { + CUDA_CHECK(cudaMemset(output + idx * outputElem, 0, sizeof(float))); } int numElem = 0; for (unsigned int i = 0; i < mYoloKernel.size(); ++i) { const auto& yolo = mYoloKernel[i]; - numElem = yolo.width*yolo.height*batchSize; + numElem = yolo.width*yolo.height*batchSize; if (numElem < mThreadCount) mThreadCount = numElem; - CalDetection<<< (yolo.width*yolo.height*batchSize + mThreadCount - 1) / mThreadCount, mThreadCount>>> - (inputs[i], output, numElem, yolo.width, yolo.height, (float *)mAnchor[i], mClassCount, outputElem); - } + //printf("Net: %d %d \n", mYoloV5NetWidth, mYoloV5NetHeight); + CalDetection << < (yolo.width*yolo.height*batchSize + mThreadCount - 1) / mThreadCount, mThreadCount >> > + (inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetWidth, mMaxOutObject, yolo.width, yolo.height, (float *)mAnchor[i], mClassCount, outputElem); + } } @@ -243,22 +257,81 @@ namespace nvinfer1 const char* YoloPluginCreator::getPluginName() const { - return "YoloLayer_TRT"; + return "YoloLayer_TRT"; } const char* YoloPluginCreator::getPluginVersion() const { - return "1"; + return "1"; } const PluginFieldCollection* YoloPluginCreator::getFieldNames() { - return &mFC; + return &mFC; } IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) { - YoloLayerPlugin* obj = new YoloLayerPlugin(); + int ClassCount; + int YoloV5NetWidth; + int YoloV5NetHeight; + int MaxOutObject; + std::vector vYoloKernel; + + int* floatdata0; + int* floatdata1; + int* floatdata2; + int* floatdata3; + const PluginField* fields = fc->fields; + for (int i = 0; i < fc->nbFields; i++) { + + if (strcmp(fields[i].name, "netdata") == 0) { + assert(fields[i].type == PluginFieldType::kFLOAT32); + floatdata0 = (int*)(fields[i].data); + ClassCount = floatdata0[0]; + YoloV5NetWidth = floatdata0[1]; + YoloV5NetHeight = floatdata0[2]; + MaxOutObject = floatdata0[3]; + } + if (strcmp(fields[i].name, "yolodata1") == 0) { + assert(fields[i].type == PluginFieldType::kFLOAT32); + floatdata1 = (int*)(fields[i].data); + YoloKernel kernel; + kernel.width = floatdata1[0]; + kernel.height = floatdata1[1]; + for (int j = 0; j < fields[i].length - 2; j++) + { + kernel.anchors[j] = floatdata1[j + 2]; + } + vYoloKernel.push_back(kernel); + } + if (strcmp(fields[i].name, "yolodata2") == 0) { + assert(fields[i].type == PluginFieldType::kFLOAT32); + floatdata2 = (int*)(fields[i].data); + YoloKernel kernel; + kernel.width = floatdata2[0]; + kernel.height = floatdata2[1]; + for (int j = 0; j < fields[i].length - 2; j++) + { + kernel.anchors[j] = floatdata2[j + 2]; + } + vYoloKernel.push_back(kernel); + } + if (strcmp(fields[i].name, "yolodata3") == 0) { + assert(fields[i].type == PluginFieldType::kFLOAT32); + floatdata3 = (int*)(fields[i].data); + YoloKernel kernel; + kernel.width = floatdata3[0]; + kernel.height = floatdata3[1]; + for (int j = 0; j < fields[i].length - 2; j++) + { + kernel.anchors[j] = floatdata3[j + 2]; + } + vYoloKernel.push_back(kernel); + } + } + std::reverse(vYoloKernel.begin(), vYoloKernel.end()); + YoloLayerPlugin* obj = new YoloLayerPlugin(ClassCount, YoloV5NetWidth, YoloV5NetHeight, MaxOutObject, vYoloKernel); obj->setPluginNamespace(mNamespace.c_str()); return obj; } @@ -271,5 +344,5 @@ namespace nvinfer1 obj->setPluginNamespace(mNamespace.c_str()); return obj; } - + REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); } diff --git a/yolov5/yololayer.h b/yolov5/yololayer.h index 91116cd..a6d9a89 100644 --- a/yolov5/yololayer.h +++ b/yolov5/yololayer.h @@ -9,36 +9,19 @@ namespace Yolo { static constexpr int CHECK_COUNT = 3; static constexpr float IGNORE_THRESH = 0.1f; + struct YoloKernel + { + int width; + int height; + float anchors[CHECK_COUNT * 2]; + }; static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000; static constexpr int CLASS_NUM = 80; static constexpr int INPUT_H = 608; static constexpr int INPUT_W = 608; - struct YoloKernel - { - int width; - int height; - float anchors[CHECK_COUNT*2]; - }; - - static constexpr YoloKernel yolo1 = { - INPUT_W / 32, - INPUT_H / 32, - {116,90, 156,198, 373,326} - }; - static constexpr YoloKernel yolo2 = { - INPUT_W / 16, - INPUT_H / 16, - {30,61, 62,45, 59,119} - }; - static constexpr YoloKernel yolo3 = { - INPUT_W / 8, - INPUT_H / 8, - {10,13, 16,30, 33,23} - }; - static constexpr int LOCATIONS = 4; - struct alignas(float) Detection{ + struct alignas(float) Detection { //center_x center_y w h float bbox[LOCATIONS]; float conf; // bbox_conf * cls_conf @@ -48,105 +31,106 @@ namespace Yolo namespace nvinfer1 { - class YoloLayerPlugin: public IPluginV2IOExt + class YoloLayerPlugin : public IPluginV2IOExt { - public: - explicit YoloLayerPlugin(); - YoloLayerPlugin(const void* data, size_t length); + public: + YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector& vYoloKernel); + YoloLayerPlugin(const void* data, size_t length); + ~YoloLayerPlugin(); - ~YoloLayerPlugin(); + int getNbOutputs() const override + { + return 1; + } - int getNbOutputs() const override - { - return 1; - } + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override; - Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override; + int initialize() override; - int initialize() override; + virtual void terminate() override {}; - virtual void terminate() override {}; + virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0; } - virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0;} + virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override; - virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override; + virtual size_t getSerializationSize() const override; - virtual size_t getSerializationSize() const override; + virtual void serialize(void* buffer) const override; - virtual void serialize(void* buffer) const override; + bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override { + return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; + } - bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override { - return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; - } + const char* getPluginType() const override; - const char* getPluginType() const override; + const char* getPluginVersion() const override; - const char* getPluginVersion() const override; + void destroy() override; - void destroy() override; + IPluginV2IOExt* clone() const override; - IPluginV2IOExt* clone() const override; + void setPluginNamespace(const char* pluginNamespace) override; - void setPluginNamespace(const char* pluginNamespace) override; + const char* getPluginNamespace() const override; - const char* getPluginNamespace() const override; + DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override; - DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override; + bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override; - bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override; + bool canBroadcastInputAcrossBatch(int inputIndex) const override; - bool canBroadcastInputAcrossBatch(int inputIndex) const override; + void attachToContext( + cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override; - void attachToContext( - cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override; + void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override; - void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override; + void detachFromContext() override; - void detachFromContext() override; - - private: - void forwardGpu(const float *const * inputs,float * output, cudaStream_t stream,int batchSize = 1); - int mClassCount; - int mKernelCount; - std::vector mYoloKernel; - int mThreadCount = 256; - void** mAnchor; - const char* mPluginNamespace; + private: + void forwardGpu(const float *const * inputs, float * output, cudaStream_t stream, int batchSize = 1); + int mThreadCount = 256; + const char* mPluginNamespace; + int mKernelCount; + int mClassCount; + int mYoloV5NetWidth; + int mYoloV5NetHeight; + int mMaxOutObject; + std::vector mYoloKernel; + void** mAnchor; }; class YoloPluginCreator : public IPluginCreator { - public: - YoloPluginCreator(); + public: + YoloPluginCreator(); - ~YoloPluginCreator() override = default; + ~YoloPluginCreator() override = default; - const char* getPluginName() const override; + const char* getPluginName() const override; - const char* getPluginVersion() const override; + const char* getPluginVersion() const override; - const PluginFieldCollection* getFieldNames() override; + const PluginFieldCollection* getFieldNames() override; - IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override; + IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override; - IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override; + IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override; - void setPluginNamespace(const char* libNamespace) override - { - mNamespace = libNamespace; - } + void setPluginNamespace(const char* libNamespace) override + { + mNamespace = libNamespace; + } - const char* getPluginNamespace() const override - { - return mNamespace.c_str(); - } + const char* getPluginNamespace() const override + { + return mNamespace.c_str(); + } - private: - std::string mNamespace; - static PluginFieldCollection mFC; - static std::vector mPluginAttributes; + private: + std::string mNamespace; + static PluginFieldCollection mFC; + static std::vector mPluginAttributes; }; - REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); }; #endif diff --git a/yolov5/yolov5.cpp b/yolov5/yolov5.cpp index c0503b0..a39c17a 100644 --- a/yolov5/yolov5.cpp +++ b/yolov5/yolov5.cpp @@ -19,21 +19,23 @@ // stuff we know about the network and the input/output blobs static const int INPUT_H = Yolo::INPUT_H; static const int INPUT_W = Yolo::INPUT_W; +static const int CLASS_NUM = Yolo::CLASS_NUM; static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1; // we assume the yololayer outputs no more than 1000 boxes that conf >= 0.1 const char* INPUT_BLOB_NAME = "data"; const char* OUTPUT_BLOB_NAME = "prob"; static Logger gLogger; + // Creat the engine using only the API and not any parser. ICudaEngine* createEngine_s(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { INetworkDefinition* network = builder->createNetworkV2(0U); // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME - ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W}); + ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W }); assert(data); std::map weightMap = loadWeights("../yolov5s.wts"); - Weights emptywts{DataType::kFLOAT, nullptr, 0}; + Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; // yolov5 backbone auto focus0 = focus(network, weightMap, *data, 3, 32, 3, "model.0"); @@ -54,46 +56,40 @@ ICudaEngine* createEngine_s(unsigned int maxBatchSize, IBuilder* builder, IBuild for (int i = 0; i < 256 * 2 * 2; i++) { deval[i] = 1.0; } - Weights deconvwts11{DataType::kFLOAT, deval, 256 * 2 * 2}; - IDeconvolutionLayer* deconv11 = network->addDeconvolutionNd(*conv10->getOutput(0), 256, DimsHW{2, 2}, deconvwts11, emptywts); - deconv11->setStrideNd(DimsHW{2, 2}); + Weights deconvwts11{ DataType::kFLOAT, deval, 256 * 2 * 2 }; + IDeconvolutionLayer* deconv11 = network->addDeconvolutionNd(*conv10->getOutput(0), 256, DimsHW{ 2, 2 }, deconvwts11, emptywts); + deconv11->setStrideNd(DimsHW{ 2, 2 }); deconv11->setNbGroups(256); weightMap["deconv11"] = deconvwts11; - ITensor* inputTensors12[] = {deconv11->getOutput(0), bottleneck_csp6->getOutput(0)}; + ITensor* inputTensors12[] = { deconv11->getOutput(0), bottleneck_csp6->getOutput(0) }; auto cat12 = network->addConcatenation(inputTensors12, 2); auto bottleneck_csp13 = bottleneckCSP(network, weightMap, *cat12->getOutput(0), 512, 256, 1, false, 1, 0.5, "model.13"); auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), 128, 1, 1, 1, "model.14"); - Weights deconvwts15{DataType::kFLOAT, deval, 128 * 2 * 2}; - IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), 128, DimsHW{2, 2}, deconvwts15, emptywts); - deconv15->setStrideNd(DimsHW{2, 2}); + Weights deconvwts15{ DataType::kFLOAT, deval, 128 * 2 * 2 }; + IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), 128, DimsHW{ 2, 2 }, deconvwts15, emptywts); + deconv15->setStrideNd(DimsHW{ 2, 2 }); deconv15->setNbGroups(128); - //weightMap["deconv15"] = deconvwts15; - ITensor* inputTensors16[] = {deconv15->getOutput(0), bottleneck_csp4->getOutput(0)}; + ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) }; auto cat16 = network->addConcatenation(inputTensors16, 2); auto bottleneck_csp17 = bottleneckCSP(network, weightMap, *cat16->getOutput(0), 256, 128, 1, false, 1, 0.5, "model.17"); - IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); + IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), 128, 3, 2, 1, "model.18"); - ITensor* inputTensors19[] = {conv18->getOutput(0), conv14->getOutput(0)}; + ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; auto cat19 = network->addConcatenation(inputTensors19, 2); auto bottleneck_csp20 = bottleneckCSP(network, weightMap, *cat19->getOutput(0), 256, 256, 1, false, 1, 0.5, "model.20"); - IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); + IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), 256, 3, 2, 1, "model.21"); - ITensor* inputTensors22[] = {conv21->getOutput(0), conv10->getOutput(0)}; + ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; auto cat22 = network->addConcatenation(inputTensors22, 2); auto bottleneck_csp23 = bottleneckCSP(network, weightMap, *cat22->getOutput(0), 512, 512, 1, false, 1, 0.5, "model.23"); - IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - - auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); - const PluginFieldCollection* pluginData = creator->getFieldNames(); - IPluginV2 *pluginObj = creator->createPlugin("yololayer", pluginData); - ITensor* inputTensors_yolo[] = {det2->getOutput(0), det1->getOutput(0), det0->getOutput(0)}; - auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj); + IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); + auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); network->markOutput(*yolo->getOutput(0)); @@ -113,9 +109,9 @@ ICudaEngine* createEngine_s(unsigned int maxBatchSize, IBuilder* builder, IBuild // Release host memory for (auto& mem : weightMap) { - free((void*) (mem.second.values)); - } + free((void*)(mem.second.values)); + } return engine; } @@ -166,38 +162,25 @@ ICudaEngine* createEngine_m(unsigned int maxBatchSize, IBuilder* builder, IBuild ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) }; auto cat16 = network->addConcatenation(inputTensors16, 2); - auto bottleneck_csp17 = bottleneckCSP(network, weightMap, *cat16->getOutput(0), 384, 192, 2, false, 1, 0.5, "model.17"); //yolo layer 0 IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), 192, 3, 2, 1, "model.18"); - - ITensor* inputTensors19[] = {conv18->getOutput(0), conv14->getOutput(0)}; + ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = bottleneckCSP(network, weightMap, *cat19->getOutput(0), 384, 384, 2, false, 1, 0.5, "model.20"); //yolo layer 1 IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), 384, 3, 2, 1, "model.21"); - ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = bottleneckCSP(network, weightMap, *cat22->getOutput(0), 768, 768, 2, false, 1, 0.5, "model.23"); - // yolo layer 2 IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); - const PluginFieldCollection* pluginData = creator->getFieldNames(); - IPluginV2 *pluginObj = creator->createPlugin("yololayer", pluginData); - ITensor* inputTensors_yolo[] = {det2->getOutput(0), det1->getOutput(0), det0->getOutput(0)}; - auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj); - + auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); network->markOutput(*yolo->getOutput(0)); @@ -267,39 +250,26 @@ ICudaEngine* createEngine_l(unsigned int maxBatchSize, IBuilder* builder, IBuild IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), 256, DimsHW{ 2, 2 }, deconvwts15, emptywts); deconv15->setStrideNd(DimsHW{ 2, 2 }); deconv15->setNbGroups(256); - ITensor* inputTensors16[] = {deconv15->getOutput(0), bottleneck_csp4->getOutput(0)}; + ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) }; auto cat16 = network->addConcatenation(inputTensors16, 2); auto bottleneck_csp17 = bottleneckCSP(network, weightMap, *cat16->getOutput(0), 512, 256, 3, false, 1, 0.5, "model.17"); // yolo layer 0 IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), 256, 3, 2, 1, "model.18"); - - ITensor* inputTensors19[] = {conv18->getOutput(0), conv14->getOutput(0)}; + ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = bottleneckCSP(network, weightMap, *cat19->getOutput(0), 512, 512, 3, false, 1, 0.5, "model.20"); - //yolo layer 1 IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), 512, 3, 2, 1, "model.21"); - - ITensor* inputTensors22[] = {conv21->getOutput(0), conv10->getOutput(0)}; + ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = bottleneckCSP(network, weightMap, *cat22->getOutput(0), 1024, 1024, 3, false, 1, 0.5, "model.23"); - IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); - const PluginFieldCollection* pluginData = creator->getFieldNames(); - IPluginV2 *pluginObj = creator->createPlugin("yololayer", pluginData); - ITensor* inputTensors_yolo[] = {det2->getOutput(0), det1->getOutput(0), det0->getOutput(0)}; - auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj); - + auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); network->markOutput(*yolo->getOutput(0)); @@ -377,33 +347,20 @@ ICudaEngine* createEngine_x(unsigned int maxBatchSize, IBuilder* builder, IBuild // yolo layer 0 IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), 320, 3, 2, 1, "model.18"); - ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = bottleneckCSP(network, weightMap, *cat19->getOutput(0), 640, 640, 4, false, 1, 0.5, "model.20"); - // yolo layer 1 IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), 640, 3, 2, 1, "model.21"); - ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = bottleneckCSP(network, weightMap, *cat22->getOutput(0), 1280, 1280, 4, false, 1, 0.5, "model.23"); - // yolo layer 2 IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); - const PluginFieldCollection* pluginData = creator->getFieldNames(); - IPluginV2 *pluginObj = creator->createPlugin("yololayer", pluginData); - ITensor* inputTensors_yolo[] = { det2->getOutput(0), det1->getOutput(0), det0->getOutput(0) }; - auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj); - + auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); network->markOutput(*yolo->getOutput(0)); @@ -458,12 +415,12 @@ void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffer int main(int argc, char** argv) { cudaSetDevice(DEVICE); // create a model using the API directly and serialize it to a stream - char *trtModelStream{nullptr}; - size_t size{0}; + char *trtModelStream{ nullptr }; + size_t size{ 0 }; std::string engine_name = STR2(NET); engine_name = "yolov5" + engine_name + ".engine"; if (argc == 2 && std::string(argv[1]) == "-s") { - IHostMemory* modelStream{nullptr}; + IHostMemory* modelStream{ nullptr }; APIToModel(BATCH_SIZE, &modelStream); assert(modelStream != nullptr); std::ofstream p(engine_name, std::ios::binary); @@ -474,7 +431,8 @@ int main(int argc, char** argv) { p.write(reinterpret_cast(modelStream->data()), modelStream->size()); modelStream->destroy(); return 0; - } else if (argc == 3 && std::string(argv[1]) == "-d") { + } + else if (argc == 3 && std::string(argv[1]) == "-d") { std::ifstream file(engine_name, std::ios::binary); if (file.good()) { file.seekg(0, file.end); @@ -485,7 +443,8 @@ int main(int argc, char** argv) { file.read(trtModelStream, size); file.close(); } - } else { + } + else { std::cerr << "arguments not right!" << std::endl; std::cerr << "./yolov5 -s // serialize model to plan file" << std::endl; std::cerr << "./yolov5 -d ../samples // deserialize plan file and run inference" << std::endl;