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

cleanup, update readme

Browse Source
This commit is contained in:
wang-xinyu 2020-05-03 08:18:47 +08:00
parent 3467ca224b
commit 2fd37375bc
5 changed files with 20 additions and 120 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
yolov4/README.md
View File

@ -43,3 +43,16 @@ sudo ./yolov4 -d ../../yolov3-spp/samples // deserialize plan file and run infe
<p align="center">
<img src="https://user-images.githubusercontent.com/15235574/80863730-cfffc500-8cb0-11ea-810e-94d693e71d80.jpg">
</p>
## 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

1
yolov4/mish.cu
View File

@ -32,7 +32,6 @@ namespace nvinfer1
int MishPlugin::initialize()
{
printf("input size : %d \n", input_size_);
return 0;
}

105
yolov4/yololayer.cu
View File

@ -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 <Detection> 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;
};

6
yolov4/yololayer.h
View File

@ -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<Yolo::YoloKernel> mYoloKernel;
int mThreadCount;
//int mDetNum;
//cpu
void* mInputBuffer {nullptr};
void* mOutputBuffer {nullptr};
};
};

15
yolov4/yolov4.cpp
View File

@ -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<Yolo::Detection>& res, float *output, float nms_thresh = 0.4) {
void nms(std::vector<Yolo::Detection>& res, float *output, float nms_thresh = NMS_THRESH) {
std::map<float, std::vector<Yolo::Detection>> 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<Yolo::Detection>());
@ -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<std::chrono::milliseconds>(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];