diff --git a/README.md b/README.md index 4d48695..0469fde 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,7 @@ Following models are implemented. |[yolov3](./yolov3)| darknet-53, weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) | |[yolov3-spp](./yolov3-spp)| darknet-53, weights and pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) | |[yolov4](./yolov4)| CSPDarknet53, weights from [AlexeyAB/darknet](https://github.com/AlexeyAB/darknet#pre-trained-models), pytorch implementation from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) | +|[yolov5](./yolov5)| yolov5-s, pytorch implementation from [ultralytics/yolov5](https://github.com/ultralytics/yolov5) | |[retinaface](./retinaface)| resnet-50, weights from [biubug6/Pytorch_Retinaface](https://github.com/biubug6/Pytorch_Retinaface) | |[arcface](./arcface)| LResNet50E-IR, weights from [deepinsight/insightface](https://github.com/deepinsight/insightface) | |[retinafaceAntiCov](./retinafaceAntiCov)| mobilenet0.25, weights from [deepinsight/insightface](https://github.com/deepinsight/insightface), retinaface anti-COVID-19, detect face and mask attribute | @@ -85,6 +86,9 @@ Some tricky operations encountered in these models, already solved, but might ha | YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 1 | FP16 | 608x608 | 35.7 | | YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 4 | FP16 | 608x608 | 40.9 | | YOLOv4(CSPDarknet53) | Xeon E5-2620/GTX1080 | 8 | FP16 | 608x608 | 41.3 | +| YOLOv5-s | Xeon E5-2620/GTX1080 | 1 | FP16 | 608x608 | 167 | +| YOLOv5-s | Xeon E5-2620/GTX1080 | 4 | FP16 | 608x608 | 182 | +| YOLOv5-s | Xeon E5-2620/GTX1080 | 8 | FP16 | 608x608 | 186 | | RetinaFace(resnet50) | TX2 | 1 | FP16 | 384x640 | 15 | | RetinaFace(resnet50) | Xeon E5-2620/GTX1080 | 1 | FP32 | 928x1600 | 15 | | ArcFace(LResNet50E-IR) | Xeon E5-2620/GTX1080 | 1 | FP32 | 112x112 | 333 | diff --git a/yolov5/README.md b/yolov5/README.md index eb900b7..cdff942 100644 --- a/yolov5/README.md +++ b/yolov5/README.md @@ -45,6 +45,7 @@ sudo ./yolov5s -d ../samples // deserialize plan file and run inference, the im - GPU id can be selected by the macro in yolov5s.cpp - NMS thresh in yolov5s.cpp - BBox confidence thresh in yolov5s.cpp +- Batch size in yolov5s.cpp ## More Information diff --git a/yolov5/yololayer.cu b/yolov5/yololayer.cu index aef5d60..86a67d6 100644 --- a/yolov5/yololayer.cu +++ b/yolov5/yololayer.cu @@ -15,6 +15,15 @@ namespace nvinfer1 mYoloKernel.push_back(yolo3); mKernelCount = mYoloKernel.size(); + + CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); + size_t AnchorLen = sizeof(float)* CHECK_COUNT*2; + for(int ii = 0; ii < mKernelCount; ii ++) + { + CUDA_CHECK(cudaMalloc(&mAnchor[ii],AnchorLen)); + const auto& yolo = mYoloKernel[ii]; + CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); + } } YoloLayerPlugin::~YoloLayerPlugin() @@ -34,6 +43,15 @@ namespace nvinfer1 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 ++) + { + CUDA_CHECK(cudaMalloc(&mAnchor[ii],AnchorLen)); + const auto& yolo = mYoloKernel[ii]; + CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); + } + assert(d == a + length); } @@ -162,7 +180,7 @@ namespace nvinfer1 float *res_count = output + bnIdx*outputElem; int count = (int)atomicAdd(res_count, 1); if (count >= MAX_OUTPUT_BBOX_COUNT) return; - char* data = (char * )res_count + sizeof(float) + count*sizeof(Detection); + char* data = (char *)res_count + sizeof(float) + count * sizeof(Detection); Detection* det = (Detection*)(data); int row = idx / yoloWidth; @@ -181,9 +199,6 @@ namespace nvinfer1 } void YoloLayerPlugin::forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize) { - void* devAnchor; - size_t AnchorLen = sizeof(float)* CHECK_COUNT*2; - CUDA_CHECK(cudaMalloc(&devAnchor,AnchorLen)); int outputElem = 1 + MAX_OUTPUT_BBOX_COUNT * sizeof(Detection) / sizeof(float); @@ -191,28 +206,22 @@ namespace nvinfer1 CUDA_CHECK(cudaMemset(output + idx*outputElem, 0, sizeof(float))); } int numElem = 0; - for (unsigned int i = 0;i< mYoloKernel.size();++i) + for (unsigned int i = 0; i < mYoloKernel.size(); ++i) { const auto& yolo = mYoloKernel[i]; numElem = yolo.width*yolo.height*batchSize; if (numElem < mThreadCount) mThreadCount = numElem; - CUDA_CHECK(cudaMemcpy(devAnchor, yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); CalDetection<<< (yolo.width*yolo.height*batchSize + mThreadCount - 1) / mThreadCount, mThreadCount>>> - (inputs[i],output, numElem, yolo.width, yolo.height, (float *)devAnchor, mClassCount ,outputElem); + (inputs[i], output, numElem, yolo.width, yolo.height, (float *)mAnchor[i], mClassCount, outputElem); } - CUDA_CHECK(cudaFree(devAnchor)); } int YoloLayerPlugin::enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) { - //assert(batchSize == 1); - //GPU - //CUDA_CHECK(cudaStreamSynchronize(stream)); forwardGpu((const float *const *)inputs, (float*)outputs[0], stream, batchSize); - return 0; } diff --git a/yolov5/yololayer.h b/yolov5/yololayer.h index cc642af..c3a1357 100644 --- a/yolov5/yololayer.h +++ b/yolov5/yololayer.h @@ -110,6 +110,7 @@ namespace nvinfer1 int mKernelCount; std::vector mYoloKernel; int mThreadCount = 256; + void** mAnchor; const char* mPluginNamespace; }; diff --git a/yolov5/yolov5s.cpp b/yolov5/yolov5s.cpp index c9dbd48..cc2f951 100644 --- a/yolov5/yolov5s.cpp +++ b/yolov5/yolov5s.cpp @@ -8,6 +8,7 @@ #define DEVICE 0 // GPU id #define NMS_THRESH 0.5 #define CONF_THRESH 0.4 +#define BATCH_SIZE 1 // stuff we know about the network and the input/output blobs static const int INPUT_H = Yolo::INPUT_H; @@ -159,7 +160,7 @@ int main(int argc, char** argv) { if (argc == 2 && std::string(argv[1]) == "-s") { IHostMemory* modelStream{nullptr}; - APIToModel(1, &modelStream); + APIToModel(BATCH_SIZE, &modelStream); assert(modelStream != nullptr); std::ofstream p("yolov5s.engine", std::ios::binary); if (!p) { @@ -194,10 +195,10 @@ int main(int argc, char** argv) { } // prepare input data --------------------------- - float data[3 * INPUT_H * INPUT_W]; + static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W]; //for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++) // data[i] = 1.0; - static float prob[OUTPUT_SIZE]; + static float prob[BATCH_SIZE * OUTPUT_SIZE]; IRuntime* runtime = createInferRuntime(gLogger); assert(runtime != nullptr); ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); @@ -207,40 +208,42 @@ int main(int argc, char** argv) { delete[] trtModelStream; int fcount = 0; - for (auto f: file_names) { + for (int f = 0; f < (int)file_names.size(); f++) { fcount++; - std::cout << fcount << " " << f << std::endl; - cv::Mat img = cv::imread(std::string(argv[2]) + "/" + f); - if (img.empty()) continue; - cv::Mat pr_img = preprocess_img(img); - for (int i = 0; i < INPUT_H * INPUT_W; i++) { - data[i] = pr_img.at(i)[2] / 255.0; - data[i + INPUT_H * INPUT_W] = pr_img.at(i)[1] / 255.0; - data[i + 2 * INPUT_H * INPUT_W] = pr_img.at(i)[0] / 255.0; + if (fcount < BATCH_SIZE && f + 1 != (int)file_names.size()) continue; + for (int b = 0; b < fcount; b++) { + cv::Mat img = cv::imread(std::string(argv[2]) + "/" + file_names[f - fcount + 1 + b]); + if (img.empty()) continue; + cv::Mat pr_img = preprocess_img(img); + for (int i = 0; i < INPUT_H * INPUT_W; i++) { + data[b * 3 * INPUT_H * INPUT_W + i] = pr_img.at(i)[2] / 255.0; + data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = pr_img.at(i)[1] / 255.0; + data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = pr_img.at(i)[0] / 255.0; + } } // Run inference auto start = std::chrono::system_clock::now(); - doInference(*context, data, prob, 1); + doInference(*context, data, prob, BATCH_SIZE); auto end = std::chrono::system_clock::now(); std::cout << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; - std::vector res; - nms(res, prob, CONF_THRESH, NMS_THRESH); - for (int i=0; i<20; i++) { - std::cout << prob[i] << ","; + std::vector> batch_res(fcount); + for (int b = 0; b < fcount; b++) { + auto& res = batch_res[b]; + nms(res, &prob[b * OUTPUT_SIZE], CONF_THRESH, NMS_THRESH); } - std::cout << res.size() << std::endl; - for (size_t j = 0; j < res.size(); j++) { - float *p = (float*)&res[j]; - for (size_t k = 0; k < 7; k++) { - std::cout << p[k] << ", "; + for (int b = 0; b < fcount; b++) { + auto& res = batch_res[b]; + //std::cout << res.size() << std::endl; + cv::Mat img = cv::imread(std::string(argv[2]) + "/" + file_names[f - fcount + 1 + b]); + for (size_t j = 0; j < res.size(); j++) { + cv::Rect r = get_rect(img, res[j].bbox); + cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2); + cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2); } - std::cout << std::endl; - cv::Rect r = get_rect(img, res[j].bbox); - cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2); - cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2); + cv::imwrite("_" + file_names[f - fcount + 1 + b], img); } - cv::imwrite("_" + f, img); + fcount = 0; } // Destroy the engine