diff --git a/yolov4/README.md b/yolov4/README.md index b70c18b..3b8dca4 100644 --- a/yolov4/README.md +++ b/yolov4/README.md @@ -43,3 +43,16 @@ sudo ./yolov4 -d ../../yolov3-spp/samples // deserialize plan file and run infe

+ +## Config + +- Input shape defined in yololayer.h +- Number of classes defined in yololayer.h +- FP16/FP32 can be selected by the macro in yolov4.cpp +- GPU id can be selected by the macro in yolov4.cpp +- NMS thresh in yolov4.cpp +- BBox confidence thresh in yolov4.cpp + +## More Information + +See the [readme](../README.md) in home page diff --git a/yolov4/mish.cu b/yolov4/mish.cu index 4637bde..cbef705 100644 --- a/yolov4/mish.cu +++ b/yolov4/mish.cu @@ -32,7 +32,6 @@ namespace nvinfer1 int MishPlugin::initialize() { - printf("input size : %d \n", input_size_); return 0; } diff --git a/yolov4/yololayer.cu b/yolov4/yololayer.cu index 7db4107..3804b45 100644 --- a/yolov4/yololayer.cu +++ b/yolov4/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/yolov4/yololayer.h b/yolov4/yololayer.h index 51f4d77..8ffcc19 100644 --- a/yolov4/yololayer.h +++ b/yolov4/yololayer.h @@ -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/yolov4/yolov4.cpp b/yolov4/yolov4.cpp index 87a8779..f8ad3da 100644 --- a/yolov4/yolov4.cpp +++ b/yolov4/yolov4.cpp @@ -16,6 +16,8 @@ //#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; @@ -94,10 +96,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()); @@ -434,7 +436,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType auto l135 = convBnLeaky(network, weightMap, *l134->getOutput(0), 256, 3, 1, 1, 135); auto l136 = convBnLeaky(network, weightMap, *l135->getOutput(0), 128, 1, 1, 0, 136); auto l137 = convBnLeaky(network, weightMap, *l136->getOutput(0), 256, 3, 1, 1, 137); - IConvolutionLayer* conv138 = network->addConvolution(*l137->getOutput(0), 255, DimsHW{1, 1}, weightMap["module_list.138.Conv2d.weight"], weightMap["module_list.138.Conv2d.bias"]); + IConvolutionLayer* conv138 = network->addConvolution(*l137->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["module_list.138.Conv2d.weight"], weightMap["module_list.138.Conv2d.bias"]); assert(conv138); // 139 is yolo layer @@ -450,7 +452,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType auto l146 = convBnLeaky(network, weightMap, *l145->getOutput(0), 512, 3, 1, 1, 146); auto l147 = convBnLeaky(network, weightMap, *l146->getOutput(0), 256, 1, 1, 0, 147); auto l148 = convBnLeaky(network, weightMap, *l147->getOutput(0), 512, 3, 1, 1, 148); - IConvolutionLayer* conv149 = network->addConvolution(*l148->getOutput(0), 255, DimsHW{1, 1}, weightMap["module_list.149.Conv2d.weight"], weightMap["module_list.149.Conv2d.bias"]); + IConvolutionLayer* conv149 = network->addConvolution(*l148->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["module_list.149.Conv2d.weight"], weightMap["module_list.149.Conv2d.bias"]); assert(conv149); // 150 is yolo layer @@ -466,7 +468,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType auto l157 = convBnLeaky(network, weightMap, *l156->getOutput(0), 1024, 3, 1, 1, 157); auto l158 = convBnLeaky(network, weightMap, *l157->getOutput(0), 512, 1, 1, 0, 158); auto l159 = convBnLeaky(network, weightMap, *l158->getOutput(0), 1024, 3, 1, 1, 159); - IConvolutionLayer* conv160 = network->addConvolution(*l159->getOutput(0), 255, DimsHW{1, 1}, weightMap["module_list.160.Conv2d.weight"], weightMap["module_list.160.Conv2d.bias"]); + IConvolutionLayer* conv160 = network->addConvolution(*l159->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["module_list.160.Conv2d.weight"], weightMap["module_list.160.Conv2d.bias"]); assert(conv160); // 161 is yolo layer @@ -647,9 +649,6 @@ int main(int argc, char** argv) { nms(res, prob); auto end = std::chrono::system_clock::now(); std::cout << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; - for (int i=0; i<20; i++) { - std::cout << prob[i] << ","; - } std::cout << res.size() << std::endl; for (size_t j = 0; j < res.size(); j++) { float *p = (float*)&res[j];