diff --git a/yolov3-spp/README.md b/yolov3-spp/README.md index e462b7b..3c28a0a 100644 --- a/yolov3-spp/README.md +++ b/yolov3-spp/README.md @@ -1,5 +1,7 @@ # yolov3-spp +yolov4 is [here](../yolov4). + The Pytorch implementation is [ultralytics/yolov3](https://github.com/ultralytics/yolov3). It provides two trained weights of yolov3-spp, `yolov3-spp.pt` and `yolov3-spp-ultralytics.pt`(originally named `ultralytics68.pt`). Following tricks are used in this yolov3-spp: @@ -43,5 +45,16 @@ sudo ./yolov3-spp -d ../samples // deserialize plan file and run inference, the

+## Config +- Input shape defined in yololayer.h +- Number of classes defined in yololayer.h +- FP16/FP32 can be selected by the macro in yolov3-spp.cpp +- GPU id can be selected by the macro in yolov3-spp.cpp +- NMS thresh in yolov3-spp.cpp +- BBox confidence thresh in yolov3-spp.cpp + +## More Information + +See the [readme](../README.md) in home page diff --git a/yolov3-spp/yololayer.cu b/yolov3-spp/yololayer.cu index 7db4107..3804b45 100644 --- a/yolov3-spp/yololayer.cu +++ b/yolov3-spp/yololayer.cu @@ -17,11 +17,6 @@ namespace nvinfer1 YoloLayerPlugin::~YoloLayerPlugin() { - if(mInputBuffer) - CUDA_CHECK(cudaFreeHost(mInputBuffer)); - - if(mOutputBuffer) - CUDA_CHECK(cudaFreeHost(mOutputBuffer)); } // create the plugin at runtime from a byte stream @@ -64,12 +59,10 @@ namespace nvinfer1 int totalCount = 0; for(const auto& yolo : mYoloKernel) totalCount += (LOCATIONS + 1) * yolo.width*yolo.height * CHECK_COUNT; - CUDA_CHECK(cudaHostAlloc(&mInputBuffer, totalCount * sizeof(float), cudaHostAllocDefault)); totalCount = 0;//detection count for(const auto& yolo : mYoloKernel) totalCount += yolo.width*yolo.height * CHECK_COUNT; - CUDA_CHECK(cudaHostAlloc(&mOutputBuffer, sizeof(float) + totalCount * sizeof(Detection), cudaHostAllocDefault)); return 0; } @@ -83,102 +76,6 @@ namespace nvinfer1 return Dims3(totalCount + 1, 1, 1); } - /*void YoloLayerPlugin::forwardCpu(const float*const * inputs, float* outputs, cudaStream_t stream,int batchSize) - { - auto Logist = [=](float data){ - return 1./(1. + exp(-data)); - }; - - int totalOutputCount = 0; - int i = 0; - int totalCount = 0; - for(const auto& yolo : mYoloKernel) - { - totalOutputCount += yolo.width*yolo.height * CHECK_COUNT * sizeof(Detection) / sizeof(float); - totalCount += (LOCATIONS + 1 + mClassCount) * yolo.width*yolo.height * CHECK_COUNT; - ++ i; - } - - for (int idx = 0; idx < batchSize;idx++) - { - i = 0; - float* inputData = (float *)mInputBuffer;// + idx *totalCount; //if create more batch size - for(const auto& yolo : mYoloKernel) - { - int size = (LOCATIONS + 1 + mClassCount) * yolo.width*yolo.height * CHECK_COUNT; - CUDA_CHECK(cudaMemcpyAsync(inputData, (float *)inputs[i] + idx * size, size * sizeof(float), cudaMemcpyDeviceToHost, stream)); - inputData += size; - ++ i; - } - - CUDA_CHECK(cudaStreamSynchronize(stream)); - - inputData = (float *)mInputBuffer ;//+ idx *totalCount; //if create more batch size - std::vector result; - for (const auto& yolo : mYoloKernel) - { - int stride = yolo.width*yolo.height; - for (int j = 0;j < stride ;++j) - { - for (int k = 0;k < CHECK_COUNT; ++k ) - { - int beginIdx = (LOCATIONS + 1 + mClassCount)* stride *k + j; - int objIndex = beginIdx + LOCATIONS*stride; - - //check obj - float objProb = Logist(inputData[objIndex]); - if(objProb <= IGNORE_THRESH) - continue; - - //classes - int classId = -1; - float maxProb = IGNORE_THRESH; - for (int c = 0;c< mClassCount;++c){ - float cProb = Logist(inputData[beginIdx + (5 + c) * stride]) * objProb; - if(cProb > maxProb){ - maxProb = cProb; - classId = c; - } - } - - if(classId >= 0) { - Detection det; - int row = j / yolo.width; - int cols = j % yolo.width; - - //Location - det.bbox[0] = (cols + Logist(inputData[beginIdx]))/ yolo.width; - det.bbox[1] = (row + Logist(inputData[beginIdx+stride]))/ yolo.height; - det.bbox[2] = exp(inputData[beginIdx+2*stride]) * yolo.anchors[2*k]; - det.bbox[3] = exp(inputData[beginIdx+3*stride]) * yolo.anchors[2*k + 1]; - //det.classId = classId; - det.prob = maxProb; - - result.emplace_back(det); - } - } - } - - inputData += (LOCATIONS + 1 + mClassCount) * stride * CHECK_COUNT; - } - - - int detCount =result.size(); - auto data = (float *)mOutputBuffer;// + idx*(totalOutputCount + 1); //if create more batch size - float * begin = data; - //copy count; - data[0] = (float)detCount; - data++; - //copy result - memcpy(data,result.data(),result.size()*sizeof(Detection)); - - //(count + det result) - CUDA_CHECK(cudaMemcpyAsync(outputs, begin,sizeof(float) + result.size()*sizeof(Detection), cudaMemcpyHostToDevice, stream)); - - outputs += totalOutputCount + 1; - } - };*/ - __device__ float Logist(float data){ return 1./(1. + exp(-data)); }; __global__ void CalDetection(const float *input, float *output,int noElements, @@ -263,8 +160,6 @@ namespace nvinfer1 //CUDA_CHECK(cudaStreamSynchronize(stream)); forwardGpu((const float *const *)inputs,(float *)outputs[0],stream,batchSize); - //CPU - //forwardCpu((const float *const *)inputs,(float *)outputs[0],stream,batchSize); return 0; }; diff --git a/yolov3-spp/yololayer.h b/yolov3-spp/yololayer.h index 406c954..10ec336 100644 --- a/yolov3-spp/yololayer.h +++ b/yolov3-spp/yololayer.h @@ -15,8 +15,8 @@ namespace Yolo static constexpr int CHECK_COUNT = 3; static constexpr float IGNORE_THRESH = 0.1f; static constexpr int CLASS_NUM = 80; - static constexpr int INPUT_H = 608; - static constexpr int INPUT_W = 608; + static constexpr int INPUT_H = 256; + static constexpr int INPUT_W = 416; struct YoloKernel { @@ -89,18 +89,12 @@ namespace nvinfer1 void forwardGpu(const float *const * inputs,float * output, cudaStream_t stream,int batchSize = 1); - void forwardCpu(const float *const * inputs,float * output, cudaStream_t stream,int batchSize = 1); - private: int mClassCount; int mKernelCount; std::vector mYoloKernel; int mThreadCount; //int mDetNum; - - //cpu - void* mInputBuffer {nullptr}; - void* mOutputBuffer {nullptr}; }; }; diff --git a/yolov3-spp/yolov3-spp.cpp b/yolov3-spp/yolov3-spp.cpp index e007cfb..a18fb3f 100644 --- a/yolov3-spp/yolov3-spp.cpp +++ b/yolov3-spp/yolov3-spp.cpp @@ -13,8 +13,10 @@ #include #include -#define USE_FP16 // comment out this if want to use FP32 +//#define USE_FP16 // comment out this if want to use FP32 #define DEVICE 0 // GPU id +#define NMS_THRESH 0.4 +#define BBOX_CONF_THRESH 0.5 using namespace nvinfer1; @@ -93,10 +95,10 @@ bool cmp(Yolo::Detection& a, Yolo::Detection& b) { return a.det_confidence > b.det_confidence; } -void nms(std::vector& res, float *output, float nms_thresh = 0.4) { +void nms(std::vector& res, float *output, float nms_thresh = NMS_THRESH) { std::map> m; for (int i = 0; i < output[0] && i < 1000; i++) { - if (output[1 + 7 * i + 4] <= 0.5) continue; + if (output[1 + 7 * i + 4] <= BBOX_CONF_THRESH) continue; Yolo::Detection det; memcpy(&det, &output[1 + 7 * i], 7 * sizeof(float)); if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); @@ -320,7 +322,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType auto lr85 = convBnLeaky(network, weightMap, *lr84->getOutput(0), 1024, 3, 1, 1, 85); auto lr86 = convBnLeaky(network, weightMap, *lr85->getOutput(0), 512, 1, 1, 0, 86); auto lr87 = convBnLeaky(network, weightMap, *lr86->getOutput(0), 1024, 3, 1, 1, 87); - IConvolutionLayer* conv88 = network->addConvolution(*lr87->getOutput(0), 255, DimsHW{1, 1}, weightMap["module_list.88.Conv2d.weight"], weightMap["module_list.88.Conv2d.bias"]); + IConvolutionLayer* conv88 = network->addConvolution(*lr87->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["module_list.88.Conv2d.weight"], weightMap["module_list.88.Conv2d.bias"]); assert(conv88); auto lr91 = convBnLeaky(network, weightMap, *lr86->getOutput(0), 256, 1, 1, 0, 91); @@ -343,7 +345,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType auto lr97 = convBnLeaky(network, weightMap, *lr96->getOutput(0), 512, 3, 1, 1, 97); auto lr98 = convBnLeaky(network, weightMap, *lr97->getOutput(0), 256, 1, 1, 0, 98); auto lr99 = convBnLeaky(network, weightMap, *lr98->getOutput(0), 512, 3, 1, 1, 99); - IConvolutionLayer* conv100 = network->addConvolution(*lr99->getOutput(0), 255, DimsHW{1, 1}, weightMap["module_list.100.Conv2d.weight"], weightMap["module_list.100.Conv2d.bias"]); + IConvolutionLayer* conv100 = network->addConvolution(*lr99->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["module_list.100.Conv2d.weight"], weightMap["module_list.100.Conv2d.bias"]); assert(conv100); auto lr103 = convBnLeaky(network, weightMap, *lr98->getOutput(0), 128, 1, 1, 0, 103); Weights deconvwts104{DataType::kFLOAT, deval, 128 * 2 * 2}; @@ -359,7 +361,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType auto lr109 = convBnLeaky(network, weightMap, *lr108->getOutput(0), 256, 3, 1, 1, 109); auto lr110 = convBnLeaky(network, weightMap, *lr109->getOutput(0), 128, 1, 1, 0, 110); auto lr111 = convBnLeaky(network, weightMap, *lr110->getOutput(0), 256, 3, 1, 1, 111); - IConvolutionLayer* conv112 = network->addConvolution(*lr111->getOutput(0), 255, DimsHW{1, 1}, weightMap["module_list.112.Conv2d.weight"], weightMap["module_list.112.Conv2d.bias"]); + IConvolutionLayer* conv112 = network->addConvolution(*lr111->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["module_list.112.Conv2d.weight"], weightMap["module_list.112.Conv2d.bias"]); assert(conv112); auto yolo = new YoloLayerPlugin(); ITensor* inputTensors_yolo[] = {conv88->getOutput(0), conv100->getOutput(0), conv112->getOutput(0)};