diff --git a/README.md b/README.md index 858529c..cfdb9ec 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,8 @@ All the models are implemented in pytorch or mxnet first, and export a weights f ## News +- `19 Nov 2020`. YOLOv3-SPP supports dynamic input shape, including a dynamic plugin. +- `17 Nov 2020`. [AlfengYuan](https://github.com/AlfengYuan) added a Dockerfile. - `7 Nov 2020`. All models migrated to trt7 API, and clean up the master branch. - `29 Oct 2020`. First INT8 quantization implementation! Please check retinaface. - `23 Oct 2020`. Add a .wts model zoo for quick evaluation. @@ -66,7 +68,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/m/l/x v1.0 v2.0 v3.0, pytorch implementation from [ultralytics/yolov5](https://github.com/ultralytics/yolov5) | +|[yolov5](./yolov5)| yolov5-s/m/l/x v1.0 v2.0 v3.0 v3.1, pytorch implementation from [ultralytics/yolov5](https://github.com/ultralytics/yolov5) | |[retinaface](./retinaface)| resnet50 and mobilnet0.25, 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 | diff --git a/yolov3-spp/CMakeLists.txt b/yolov3-spp/CMakeLists.txt index 0c4be5d..fd60244 100644 --- a/yolov3-spp/CMakeLists.txt +++ b/yolov3-spp/CMakeLists.txt @@ -13,20 +13,16 @@ find_package(CUDA REQUIRED) set(CUDA_NVCC_PLAGS ${CUDA_NVCC_PLAGS};-std=c++11;-g;-G;-gencode;arch=compute_30;code=sm_30) include_directories(${PROJECT_SOURCE_DIR}/include) -if (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") - message("embed_platform on") - include_directories(/usr/local/cuda/targets/aarch64-linux/include) - link_directories(/usr/local/cuda/targets/aarch64-linux/lib) -else() - message("embed_platform off") - include_directories(/usr/local/cuda/include) - link_directories(/usr/local/cuda/lib64) -endif() - +# include and link dirs of cuda and tensorrt, you need adapt them if yours are different +# cuda +include_directories(/usr/local/cuda/include) +link_directories(/usr/local/cuda/lib64) +# tensorrt +include_directories(/usr/include/x86_64-linux-gnu/) +link_directories(/usr/lib/x86_64-linux-gnu/) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") -#cuda_add_library(leaky ${PROJECT_SOURCE_DIR}/leaky.cu) cuda_add_library(yololayer SHARED ${PROJECT_SOURCE_DIR}/yololayer.cu) target_link_libraries(yololayer nvinfer cudart) diff --git a/yolov3-spp/README.md b/yolov3-spp/README.md index 9a7342d..b701517 100644 --- a/yolov3-spp/README.md +++ b/yolov3-spp/README.md @@ -1,41 +1,48 @@ # yolov3-spp -yolov4 is [here](../yolov4). +Currently this is supporting dynamic input shape, if you want to use non-dynamic version, please checkout commit [659fd2b](https://github.com/wang-xinyu/tensorrtx/commit/659fd2b23482197b19dccf746a5a3dbff1611381). 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: +## Config -- Yololayer plugin is different from the plugin used in [this repo's yolov3](https://github.com/wang-xinyu/tensorrtx/tree/master/yolov3). In this version, three yololayer are implemented in one plugin to improve speed, codes derived from [lewes6369/TensorRT-Yolov3](https://github.com/lewes6369/TensorRT-Yolov3) -- Batchnorm layer, implemented by scale layer. +- 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 +- MIN and MAX input size defined in yolov3-spp.cpp +- Optimization width and height for IOptimizationProfile defined in yolov3-spp.cpp -## Excute: +## How to Run -``` 1. generate yolov3-spp_ultralytics68.wts from pytorch implementation with yolov3-spp.cfg and yolov3-spp-ultralytics.pt, or download .wts from model zoo +``` git clone https://github.com/wang-xinyu/tensorrtx.git git clone https://github.com/ultralytics/yolov3.git // download its weights 'yolov3-spp-ultralytics.pt' -cd yolov3 -cp ../tensorrtx/yolov3-spp/gen_wts.py . +// copy gen_wts.py from tensorrtx/yolov3-spp/ to ultralytics/yolov3/ +// go to ultralytics/yolov3/ python gen_wts.py yolov3-spp-ultralytics.pt // a file 'yolov3-spp_ultralytics68.wts' will be generated. // the master branch of yolov3 should work, if not, you can checkout 4ac60018f6e6c1e24b496485f126a660d9c793d8 +``` -2. put yolov3-spp_ultralytics68.wts into yolov3-spp, build and run +2. build tensorrtx/yolov3-spp and run -mv yolov3-spp_ultralytics68.wts ../tensorrtx/yolov3-spp/ -cd ../tensorrtx/yolov3-spp +``` +// put yolov3-spp_ultralytics68.wts into tensorrtx/yolov3-spp/ +// go to tensorrtx/yolov3-spp/ mkdir build cd build cmake .. make sudo ./yolov3-spp -s // serialize model to plan file i.e. 'yolov3-spp.engine' sudo ./yolov3-spp -d ../samples // deserialize plan file and run inference, the images in samples will be processed. +``` 3. check the images generated, as follows. _zidane.jpg and _bus.jpg -```

@@ -45,15 +52,6 @@ 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 in [home page.](https://github.com/wang-xinyu/tensorrtx) diff --git a/yolov3-spp/yololayer.cu b/yolov3-spp/yololayer.cu index 722cb27..d3a7969 100644 --- a/yolov3-spp/yololayer.cu +++ b/yolov3-spp/yololayer.cu @@ -11,12 +11,25 @@ namespace nvinfer1 mYoloKernel.push_back(yolo1); mYoloKernel.push_back(yolo2); mYoloKernel.push_back(yolo3); - mKernelCount = mYoloKernel.size(); + + CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); + size_t anchorLen = sizeof(float) * CHECK_COUNT * 2; + for (int i = 0; i < mKernelCount; i++) + { + CUDA_CHECK(cudaMalloc(&mAnchor[i], anchorLen)); + const auto& yolo = mYoloKernel[i]; + CUDA_CHECK(cudaMemcpy(mAnchor[i], yolo.anchors, anchorLen, cudaMemcpyHostToDevice)); + } } - + YoloLayerPlugin::~YoloLayerPlugin() { + for (int i = 0; i < mKernelCount; i++) + { + CUDA_CHECK(cudaFree(mAnchor[i])); + } + CUDA_CHECK(cudaFreeHost(mAnchor)); } // create the plugin at runtime from a byte stream @@ -28,11 +41,19 @@ namespace nvinfer1 read(d, mThreadCount); read(d, mKernelCount); mYoloKernel.resize(mKernelCount); - auto kernelSize = mKernelCount*sizeof(YoloKernel); - memcpy(mYoloKernel.data(),d,kernelSize); + auto kernelSize = mKernelCount * sizeof(YoloKernel); + memcpy(mYoloKernel.data(), d, kernelSize); d += kernelSize; - assert(d == a + length); + + CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); + size_t anchorLen = sizeof(float) * CHECK_COUNT * 2; + for (int i = 0; i < mKernelCount; i++) + { + CUDA_CHECK(cudaMalloc(&mAnchor[i], anchorLen)); + const auto& yolo = mYoloKernel[i]; + CUDA_CHECK(cudaMemcpy(mAnchor[i], yolo.anchors, anchorLen, cudaMemcpyHostToDevice)); + } } void YoloLayerPlugin::serialize(void* buffer) const @@ -42,29 +63,32 @@ namespace nvinfer1 write(d, mClassCount); write(d, mThreadCount); write(d, mKernelCount); - auto kernelSize = mKernelCount*sizeof(YoloKernel); - memcpy(d,mYoloKernel.data(),kernelSize); + auto kernelSize = mKernelCount * sizeof(YoloKernel); + memcpy(d,mYoloKernel.data(), kernelSize); d += kernelSize; assert(d == a + getSerializationSize()); } - + size_t YoloLayerPlugin::getSerializationSize() const - { + { return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size(); } int YoloLayerPlugin::initialize() - { + { return 0; } - - Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) + + DimsExprs YoloLayerPlugin::getOutputDimensions(int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder& exprBuilder) { //output the result to channel int totalsize = MAX_OUTPUT_BBOX_COUNT * sizeof(Detection) / sizeof(float); - - return Dims3(totalsize + 1, 1, 1); + DimsExprs de; + de.nbDims = 2; + de.d[0] = exprBuilder.constant(inputs[0].d[0]->getConstantValue()); // batchsize + de.d[1] = exprBuilder.constant(totalsize + 1); // outputsize + return de; } // Set plugin namespace @@ -84,19 +108,7 @@ namespace nvinfer1 return DataType::kFLOAT; } - // Return true if output tensor is broadcast across a batch. - bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const - { - return false; - } - - // Return true if plugin can use input that is broadcast across batch without replication. - bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const - { - return false; - } - - void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) + void YoloLayerPlugin::configurePlugin(const DynamicPluginTensorDesc* in, int nbInputs, const DynamicPluginTensorDesc* out, int nbOutputs) { } @@ -124,7 +136,7 @@ namespace nvinfer1 } // Clone the plugin - IPluginV2IOExt* YoloLayerPlugin::clone() const + IPluginV2DynamicExt* YoloLayerPlugin::clone() const { YoloLayerPlugin *p = new YoloLayerPlugin(); p->setPluginNamespace(mPluginNamespace); @@ -133,9 +145,9 @@ namespace nvinfer1 __device__ float Logist(float data){ return 1.0f / (1.0f + expf(-data)); }; - __global__ void CalDetection(const float *input, float *output,int noElements, - int yoloWidth,int yoloHeight,const float anchors[CHECK_COUNT*2],int classes,int outputElem) { - + __global__ void CalDetection(const float *input, float *output, int noElements, + int yoloWidth, int yoloHeight, int yoloStride, const float anchors[CHECK_COUNT * 2], int classes, int outputElem) { + int idx = threadIdx.x + blockDim.x * blockIdx.x; if (idx >= noElements) return; @@ -158,20 +170,20 @@ namespace nvinfer1 float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]); if (max_cls_prob < IGNORE_THRESH || box_prob < IGNORE_THRESH) continue; - float *res_count = output + bnIdx*outputElem; + 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); - Detection* det = (Detection*)(data); + char* data = (char*)res_count + sizeof(float) + count * sizeof(Detection); + Detection* det = (Detection*)(data); int row = idx / yoloWidth; int col = idx % yoloWidth; //Location - det->bbox[0] = (col + Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * INPUT_W / yoloWidth; - det->bbox[1] = (row + Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * INPUT_H / yoloHeight; - det->bbox[2] = expf(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]) * anchors[2*k]; - det->bbox[3] = expf(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]) * anchors[2*k + 1]; + det->bbox[0] = (col + Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * yoloStride; + det->bbox[1] = (row + Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * yoloStride; + det->bbox[2] = expf(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]) * anchors[2 * k]; + det->bbox[3] = expf(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]) * anchors[2 * k + 1]; det->det_confidence = box_prob; det->class_id = class_id; det->class_confidence = max_cls_prob; @@ -179,38 +191,27 @@ 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); - for(int idx = 0 ; idx < batchSize; ++idx) { - CUDA_CHECK(cudaMemset(output + idx*outputElem, 0, sizeof(float))); + CUDA_CHECK(cudaMemset(output + idx * outputElem, 0, sizeof(float))); } int numElem = 0; - for (unsigned int i = 0;i< mYoloKernel.size();++i) - { + for (size_t 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); + numElem = yolo.width * yolo.height * batchSize; + CalDetection<<<(yolo.width * yolo.height * batchSize + mThreadCount - 1) / mThreadCount, mThreadCount>>> + (inputs[i], output, numElem, yolo.width, yolo.height, yolo.stride, (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) + int YoloLayerPlugin::enqueue(const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) { - //assert(batchSize == 1); - //GPU - //CUDA_CHECK(cudaStreamSynchronize(stream)); + int batchSize = inputDesc[0].dims.d[0]; + for (size_t i = 0; i < mYoloKernel.size(); ++i) { + mYoloKernel[i].width = inputDesc[i].dims.d[3]; + mYoloKernel[i].height = inputDesc[i].dims.d[2]; + } forwardGpu((const float *const *)inputs, (float*)outputs[0], stream, batchSize); - return 0; } @@ -240,17 +241,17 @@ namespace nvinfer1 return &mFC; } - IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) + IPluginV2DynamicExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) { YoloLayerPlugin* obj = new YoloLayerPlugin(); obj->setPluginNamespace(mNamespace.c_str()); return obj; } - IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) + IPluginV2DynamicExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) { // This object will be deleted when the network is destroyed, which will - // call MishPlugin::destroy() + // call YoloLayerPlugin::destroy() YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength); obj->setPluginNamespace(mNamespace.c_str()); return obj; diff --git a/yolov3-spp/yololayer.h b/yolov3-spp/yololayer.h index 1f080cb..6262ab7 100644 --- a/yolov3-spp/yololayer.h +++ b/yolov3-spp/yololayer.h @@ -15,29 +15,31 @@ namespace Yolo static constexpr float IGNORE_THRESH = 0.1f; static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000; static constexpr int CLASS_NUM = 80; - static constexpr int INPUT_H = 608; - static constexpr int INPUT_W = 608; struct YoloKernel { int width; int height; + int stride; float anchors[CHECK_COUNT*2]; }; static constexpr YoloKernel yolo1 = { - INPUT_W / 32, - INPUT_H / 32, + -1, // dynamic width and height + -1, + 32, {116,90, 156,198, 373,326} }; static constexpr YoloKernel yolo2 = { - INPUT_W / 16, - INPUT_H / 16, + -1, + -1, + 16, {30,61, 62,45, 59,119} }; static constexpr YoloKernel yolo3 = { - INPUT_W / 8, - INPUT_H / 8, + -1, + -1, + 8, {10,13, 16,30, 33,23} }; @@ -51,10 +53,9 @@ namespace Yolo }; } - namespace nvinfer1 { - class YoloLayerPlugin: public IPluginV2IOExt + class YoloLayerPlugin: public IPluginV2DynamicExt { public: explicit YoloLayerPlugin(); @@ -67,21 +68,24 @@ namespace nvinfer1 return 1; } - Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override; + //virtual Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) final; + virtual DimsExprs getOutputDimensions(int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder& exprBuilder) override; int initialize() override; virtual void terminate() override {}; - virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0;} + //virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0;} + size_t getWorkspaceSize(const PluginTensorDesc* inputs, int nbInputs, const PluginTensorDesc* outputs, int nbOutputs) const override { return 0; } - virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override; + //virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override; + int enqueue(const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) override; virtual size_t getSerializationSize() const override; virtual void serialize(void* buffer) const override; - bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override { + bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) override { return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; } @@ -91,7 +95,7 @@ namespace nvinfer1 void destroy() override; - IPluginV2IOExt* clone() const override; + IPluginV2DynamicExt* clone() const override; void setPluginNamespace(const char* pluginNamespace) override; @@ -99,14 +103,10 @@ namespace nvinfer1 DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override; - bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override; - - bool canBroadcastInputAcrossBatch(int inputIndex) const override; - void attachToContext( cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override; - void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override; + void configurePlugin(const DynamicPluginTensorDesc* in, int nbInputs, const DynamicPluginTensorDesc* out, int nbOutputs) override; void detachFromContext() override; @@ -116,6 +116,7 @@ namespace nvinfer1 int mKernelCount; std::vector mYoloKernel; int mThreadCount = 256; + void** mAnchor; const char* mPluginNamespace; }; @@ -132,9 +133,9 @@ namespace nvinfer1 const PluginFieldCollection* getFieldNames() override; - IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override; + IPluginV2DynamicExt* createPlugin(const char* name, const PluginFieldCollection* fc) override; - IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override; + IPluginV2DynamicExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override; void setPluginNamespace(const char* libNamespace) override { diff --git a/yolov3-spp/yolov3-spp.cpp b/yolov3-spp/yolov3-spp.cpp index 4fe7c65..3df8313 100644 --- a/yolov3-spp/yolov3-spp.cpp +++ b/yolov3-spp/yolov3-spp.cpp @@ -22,7 +22,6 @@ }\ } while (0) - #define USE_FP16 // comment out this if want to use FP32 #define DEVICE 0 // GPU id #define NMS_THRESH 0.4 @@ -31,58 +30,46 @@ using namespace nvinfer1; // stuff we know about the network and the input/output blobs -static const int INPUT_H = Yolo::INPUT_H; -static const int INPUT_W = Yolo::INPUT_W; -static const int OUTPUT_SIZE = 1000 * 7 + 1; // we assume the yololayer outputs no more than 1000 boxes that conf >= 0.1 +static const int MAX_INPUT_SIZE = 608; +static const int MIN_INPUT_SIZE = 128; +static const int OPT_INPUT_W = 608; +static const int OPT_INPUT_H = 608; +static const int DET_LEN = sizeof(Yolo::Detection) / sizeof(float); +static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * DET_LEN + 1; // we limit the yololayer to output no more than MAX_OUTPUT_BBOX_COUNT bboxes const char* INPUT_BLOB_NAME = "data"; const char* OUTPUT_BLOB_NAME = "prob"; static Logger gLogger; -cv::Mat preprocess_img(cv::Mat& img) { - int w, h, x, y; - float r_w = INPUT_W / (img.cols*1.0); - float r_h = INPUT_H / (img.rows*1.0); - if (r_h > r_w) { - w = INPUT_W; - h = r_w * img.rows; - x = 0; - y = (INPUT_H - h) / 2; - } else { - w = r_h* img.cols; - h = INPUT_H; - x = (INPUT_W - w) / 2; - y = 0; - } - cv::Mat re(h, w, CV_8UC3); - cv::resize(img, re, re.size(), 0, 0, cv::INTER_CUBIC); - cv::Mat out(INPUT_H, INPUT_W, CV_8UC3, cv::Scalar(128, 128, 128)); - re.copyTo(out(cv::Rect(x, y, re.cols, re.rows))); +cv::Mat letterbox(cv::Mat& img) { + float r = std::min(MAX_INPUT_SIZE / (img.cols*1.0), MAX_INPUT_SIZE / (img.rows*1.0)); + r = std::min(r, 1.0f); + int unpad_w = r * img.cols; + int unpad_h = r * img.rows; + int dw = (MAX_INPUT_SIZE - unpad_w) % 32; + int dh = (MAX_INPUT_SIZE - unpad_h) % 32; + cv::Mat re(unpad_h, unpad_w, CV_8UC3); + cv::resize(img, re, re.size()); + cv::Mat out(unpad_h + dh, unpad_w + dw, CV_8UC3, cv::Scalar(128, 128, 128)); + re.copyTo(out(cv::Rect(dw / 2, dh / 2, re.cols, re.rows))); return out; } -cv::Rect get_rect(cv::Mat& img, float bbox[4]) { - int l, r, t, b; - float r_w = INPUT_W / (img.cols * 1.0); - float r_h = INPUT_H / (img.rows * 1.0); - if (r_h > r_w) { - l = bbox[0] - bbox[2]/2.f; - r = bbox[0] + bbox[2]/2.f; - t = bbox[1] - bbox[3]/2.f - (INPUT_H - r_w * img.rows) / 2; - b = bbox[1] + bbox[3]/2.f - (INPUT_H - r_w * img.rows) / 2; - l = l / r_w; - r = r / r_w; - t = t / r_w; - b = b / r_w; - } else { - l = bbox[0] - bbox[2]/2.f - (INPUT_W - r_h * img.cols) / 2; - r = bbox[0] + bbox[2]/2.f - (INPUT_W - r_h * img.cols) / 2; - t = bbox[1] - bbox[3]/2.f; - b = bbox[1] + bbox[3]/2.f; - l = l / r_h; - r = r / r_h; - t = t / r_h; - b = b / r_h; - } +cv::Rect get_rect(cv::Size src_shape, cv::Size pre_shape, float bbox[4]) { + float ra = std::min(MAX_INPUT_SIZE / (src_shape.width * 1.0), MAX_INPUT_SIZE / (src_shape.height * 1.0)); + ra = std::min(ra, 1.0f); + int unpad_w = ra * src_shape.width; + int unpad_h = ra * src_shape.height; + int dw = (MAX_INPUT_SIZE - unpad_w) % 32; + int dh = (MAX_INPUT_SIZE - unpad_h) % 32; + + int l = bbox[0] - bbox[2]/2.f - dw / 2; + int r = bbox[0] + bbox[2]/2.f - dw / 2; + int t = bbox[1] - bbox[3]/2.f - dh / 2; + int b = bbox[1] + bbox[3]/2.f - dh / 2; + l /= ra; + r /= ra; + t /= ra; + b /= ra; return cv::Rect(l, t, r-l, b-t); } @@ -107,10 +94,10 @@ bool cmp(const Yolo::Detection& a, const Yolo::Detection& b) { 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] <= BBOX_CONF_THRESH) continue; + for (int i = 0; i < output[0] && i < Yolo::MAX_OUTPUT_BBOX_COUNT; i++) { + if (output[1 + DET_LEN * i + 4] <= BBOX_CONF_THRESH) continue; Yolo::Detection det; - memcpy(&det, &output[1 + 7 * i], 7 * sizeof(float)); + memcpy(&det, &output[1 + DET_LEN * i], DET_LEN * sizeof(float)); if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); m[det.class_id].push_back(det); } @@ -221,10 +208,10 @@ ILayer* convBnLeaky(INetworkDefinition *network, std::map& // Creat the engine using only the API and not any parser. ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { - INetworkDefinition* network = builder->createNetworkV2(0U); + const auto explicitBatch = 1U << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); + auto network = builder->createNetworkV2(explicitBatch); - // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME - ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W}); + ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims4{1, 3, -1, -1}); assert(data); std::map weightMap = loadWeights("../yolov3-spp_ultralytics68.wts"); @@ -309,7 +296,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder auto lr75 = convBnLeaky(network, weightMap, *ew74->getOutput(0), 512, 1, 1, 0, 75); auto lr76 = convBnLeaky(network, weightMap, *lr75->getOutput(0), 1024, 3, 1, 1, 76); auto lr77 = convBnLeaky(network, weightMap, *lr76->getOutput(0), 512, 1, 1, 0, 77); - + auto pool78 = network->addPoolingNd(*lr77->getOutput(0), PoolingType::kMAX, DimsHW{5,5}); pool78->setPaddingNd(DimsHW{2, 2}); pool78->setStrideNd(DimsHW{1, 1}); @@ -341,7 +328,7 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder deconv92->setStrideNd(DimsHW{2, 2}); deconv92->setNbGroups(256); weightMap["deconv92"] = deconvwts92; - + ITensor* inputTensors[] = {deconv92->getOutput(0), ew61->getOutput(0)}; auto cat93 = network->addConcatenation(inputTensors, 2); auto lr94 = convBnLeaky(network, weightMap, *cat93->getOutput(0), 256, 1, 1, 0, 94); @@ -375,11 +362,22 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder ITensor* inputTensors_yolo[] = {conv88->getOutput(0), conv100->getOutput(0), conv112->getOutput(0)}; auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj); + auto dim = yolo->getOutput(0)->getDimensions(); + std::cout << "yololayer output shape: "; + for (int i = 0; i < dim.nbDims; i++) { + std::cout << dim.d[i] << " "; + } + std::cout << std::endl; yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); network->markOutput(*yolo->getOutput(0)); + IOptimizationProfile* profile = builder->createOptimizationProfile(); + profile->setDimensions(INPUT_BLOB_NAME, OptProfileSelector::kMIN, Dims4(1, 3, MIN_INPUT_SIZE, MIN_INPUT_SIZE)); + profile->setDimensions(INPUT_BLOB_NAME, OptProfileSelector::kOPT, Dims4(1, 3, OPT_INPUT_H, OPT_INPUT_W)); + profile->setDimensions(INPUT_BLOB_NAME, OptProfileSelector::kMAX, Dims4(1, 3, MAX_INPUT_SIZE, MAX_INPUT_SIZE)); + config->addOptimizationProfile(profile); + // Build engine - builder->setMaxBatchSize(maxBatchSize); config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB #ifdef USE_FP16 config->setFlag(BuilderFlag::kFP16); @@ -417,7 +415,7 @@ void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) { builder->destroy(); } -void doInference(IExecutionContext& context, float* input, float* output, int batchSize) { +void doInference(IExecutionContext& context, float* input, float* output, cv::Size input_shape) { const ICudaEngine& engine = context.getEngine(); // Pointers to input and output device buffers to pass to engine. @@ -429,19 +427,20 @@ void doInference(IExecutionContext& context, float* input, float* output, int ba // Note that indices are guaranteed to be less than IEngine::getNbBindings() const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME); const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME); + context.setBindingDimensions(inputIndex, Dims4(1, 3, input_shape.height, input_shape.width)); // Create GPU buffers on device - CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float))); - CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float))); + CHECK(cudaMalloc(&buffers[inputIndex], 3 * input_shape.height * input_shape.width * sizeof(float))); + CHECK(cudaMalloc(&buffers[outputIndex], OUTPUT_SIZE * sizeof(float))); // Create stream cudaStream_t stream; CHECK(cudaStreamCreate(&stream)); // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host - CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream)); - context.enqueue(batchSize, buffers, stream, nullptr); - CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); + CHECK(cudaMemcpyAsync(buffers[inputIndex], input, 3 * input_shape.height * input_shape.width * sizeof(float), cudaMemcpyHostToDevice, stream)); + context.enqueueV2(buffers, stream, nullptr); + CHECK(cudaMemcpyAsync(output, buffers[outputIndex], OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); cudaStreamSynchronize(stream); // Release stream and buffers @@ -514,10 +513,6 @@ int main(int argc, char** argv) { return -1; } - // prepare input data --------------------------- - static float data[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]; IRuntime* runtime = createInferRuntime(gLogger); assert(runtime != nullptr); @@ -526,6 +521,7 @@ int main(int argc, char** argv) { IExecutionContext* context = engine->createExecutionContext(); assert(context != nullptr); delete[] trtModelStream; + context->setOptimizationProfile(0); int fcount = 0; for (auto f: file_names) { @@ -533,31 +529,21 @@ int main(int argc, char** argv) { 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; - } + cv::Mat pr_img = letterbox(img); + std::cout << "letterbox shape: " << pr_img.cols << ", " << pr_img.rows << std::endl; + if (pr_img.cols < MIN_INPUT_SIZE || pr_img.rows < MIN_INPUT_SIZE) continue; + cv::Mat blob = cv::dnn::blobFromImage(pr_img, 1.0 / 255.0, pr_img.size(), cv::Scalar(0, 0, 0), true, false); // Run inference auto start = std::chrono::system_clock::now(); - doInference(*context, data, prob, 1); + doInference(*context, blob.ptr(0), prob, pr_img.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); - for (int i=0; i<20; i++) { - std::cout << prob[i] << ","; - } - std::cout << res.size() << std::endl; + std::cout << "num of bbox: " << 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] << ", "; - } - std::cout << std::endl; - cv::Rect r = get_rect(img, res[j].bbox); + cv::Rect r = get_rect(img.size(), pr_img.size(), 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); }