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

cleanup, update yolov3-spp readme

Browse Source
This commit is contained in:
wang-xinyu 2020-05-03 08:56:10 +08:00
parent 2fd37375bc
commit 68237d5c99
4 changed files with 23 additions and 119 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
yolov3-spp/README.md
View File

@ -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
<img src="https://user-images.githubusercontent.com/15235574/78247970-60b27c00-751e-11ea-88df-41473fed4823.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 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

105
yolov3-spp/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;
};

10
yolov3-spp/yololayer.h
View File

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

14
yolov3-spp/yolov3-spp.cpp
View File

@ -13,8 +13,10 @@
#include <opencv2/opencv.hpp>
#include <dirent.h>
#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<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>());
@ -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)};