From 8e4de7dced077a74913db5e5dc340df9605b56e8 Mon Sep 17 00:00:00 2001 From: wang-xinyu Date: Sun, 28 Feb 2021 14:30:15 +0800 Subject: [PATCH] optimize yolov5 cmdline args --- yolov5/README.md | 35 +-- yolov5/samples | 1 + yolov5/yolov5.cpp | 563 +++++++++------------------------------------- 3 files changed, 121 insertions(+), 478 deletions(-) create mode 120000 yolov5/samples diff --git a/yolov5/README.md b/yolov5/README.md index cea8781..f108423 100644 --- a/yolov5/README.md +++ b/yolov5/README.md @@ -2,19 +2,22 @@ The Pytorch implementation is [ultralytics/yolov5](https://github.com/ultralytics/yolov5). +## Different versions of yolov5 + Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1 and v4.0. -- For yolov5 v4.0, please visit [yolov5 release v4.0](https://github.com/ultralytics/yolov5/releases/tag/v4.0), and use the latest commit of this repo. -- For yolov5 v3.0 and v3.1, please visit [yolov5 release v3.0](https://github.com/ultralytics/yolov5/releases/tag/v3.0) and [yolov5 release v3.1](https://github.com/ultralytics/yolov5/releases/tag/v3.1), and checkout commit ['6d0f5cb'](https://github.com/wang-xinyu/tensorrtx/commit/6d0f5cbf4745bc00b69aad54a905383fb906f103) of this repo. -- For yolov5 v2.0, please visit [yolov5 release v2.0](https://github.com/ultralytics/yolov5/releases/tag/v2.0), and checkout commit ['5cfa444'](https://github.com/wang-xinyu/tensorrtx/commit/5cfa4445170eabaa54acd5ad7f469ef65a8763f1) of this repo. -- For yolov5 v1.0, please visit [yolov5 release v1.0](https://github.com/ultralytics/yolov5/releases/tag/v1.0), and checkout commit ['f09aa3b'](https://github.com/wang-xinyu/tensorrtx/commit/f09aa3bbebf4d4d37b6d3b32a1d39e1f2678a07b) of this repo. +- For yolov5 v4.0, download .pt from [yolov5 release v4.0](https://github.com/ultralytics/yolov5/releases/tag/v4.0), `git clone -b v4.0 https://github.com/ultralytics/yolov5.git` and `git clone https://github.com/wang-xinyu/tensorrtx.git`. +- For yolov5 v3.1, download .pt from [yolov5 release v3.1](https://github.com/ultralytics/yolov5/releases/tag/v3.1), `git clone -b v3.1 https://github.com/ultralytics/yolov5.git` and `git clone https://github.com/wang-xinyu/tensorrtx.git && git checkout 6d0f5cbf4745bc00b69aad54a905383fb906f103`. +- For yolov5 v3.0, download .pt from [yolov5 release v3.0](https://github.com/ultralytics/yolov5/releases/tag/v3.0), `git clone -b v3.0 https://github.com/ultralytics/yolov5.git` and `git clone https://github.com/wang-xinyu/tensorrtx.git && git checkout 6d0f5cbf4745bc00b69aad54a905383fb906f103`. +- For yolov5 v2.0, download .pt from [yolov5 release v2.0](https://github.com/ultralytics/yolov5/releases/tag/v2.0), `git clone -b v2.0 https://github.com/ultralytics/yolov5.git` and `git clone https://github.com/wang-xinyu/tensorrtx.git && git checkout 5cfa4445170eabaa54acd5ad7f469ef65a8763f1`. +- For yolov5 v1.0, download .pt from [yolov5 release v1.0](https://github.com/ultralytics/yolov5/releases/tag/v1.0), `git clone -b v1.0 https://github.com/ultralytics/yolov5.git` and `git clone https://github.com/wang-xinyu/tensorrtx.git && git checkout f09aa3bbebf4d4d37b6d3b32a1d39e1f2678a07b`. ## Config - Choose the model s/m/l/x by `NET` macro in yolov5.cpp - Input shape defined in yololayer.h - Number of classes defined in yololayer.h, **DO NOT FORGET TO ADAPT THIS, If using your own model** -- FP16/FP32 can be selected by the macro in yolov5.cpp +- INT8/FP16/FP32 can be selected by the macro in yolov5.cpp, **INT8 need more steps, pls follow `How to Run` first and then go the `INT8 Quantization` below** - GPU id can be selected by the macro in yolov5.cpp - NMS thresh in yolov5.cpp - BBox confidence thresh in yolov5.cpp @@ -22,12 +25,11 @@ Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1 and v4.0. ## How to Run, yolov5s as example -1. generate yolov5s.wts from pytorch with yolov5s.pt, or download .wts from model zoo +1. generate .wts from pytorch with .pt, or download .wts from model zoo ``` -git clone https://github.com/wang-xinyu/tensorrtx.git -git clone https://github.com/ultralytics/yolov5.git -// download its weights 'yolov5s.pt' +// git clone src code according to `Different versions of yolov5` above +// download https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5s.pt // copy tensorrtx/yolov5/gen_wts.py into ultralytics/yolov5 // ensure the file name is yolov5s.pt and yolov5s.wts in gen_wts.py // go to ultralytics/yolov5 @@ -40,18 +42,19 @@ python gen_wts.py ``` // put yolov5s.wts into tensorrtx/yolov5 // go to tensorrtx/yolov5 -// ensure the macro NET in yolov5.cpp is s // update CLASS_NUM in yololayer.h if your model is trained on custom dataset mkdir build cd build cmake .. make -sudo ./yolov5 -s 0.33 0.50 // serialize yolov5s model to plan file i.e. 'yolov5.engine' -sudo ./yolov5 -s 0.67 0.75 // serialize yolov5m model to plan file i.e. 'yolov5.engine' -sudo ./yolov5 -s 1.0 1.0 // serialize yolov5l model to plan file i.e. 'yolov5.engine' -sudo ./yolov5 -s 1.33 1.25 // serialize yolov5x model to plan file i.e. 'yolov5.engine' -sudo ./yolov5 -s 0.17 0.25 // serialize your own yolov5 model with depth_multiple and width_multiple which you set in '.yaml' file when training the model. i.e. depth = 0.17, width = 0.25 -sudo ./yolov5 -d ../samples // deserialize plan file and run inference, the images in samples will be processed. +sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x or c gd gw] // serialize model to plan file +sudo ./yolov5 -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. +// For example yolov5s +sudo ./yolov5 -s yolov5s.wts yolov5s.engine s +sudo ./yolov5 -d yolov5s.engine ../samples +// For example Custom model with depth_multiple=0.17, width_multiple=0.25 in yolov5.yaml +sudo ./yolov5 -s yolov5_custom.wts yolov5.engine c 0.17 0.25 +sudo ./yolov5 -d yolov5.engine ../samples ``` 3. check the images generated, as follows. _zidane.jpg and _bus.jpg diff --git a/yolov5/samples b/yolov5/samples new file mode 120000 index 0000000..3a29e7c --- /dev/null +++ b/yolov5/samples @@ -0,0 +1 @@ +../yolov3-spp/samples/ \ No newline at end of file diff --git a/yolov5/yolov5.cpp b/yolov5/yolov5.cpp index 04ba668..9c95206 100644 --- a/yolov5/yolov5.cpp +++ b/yolov5/yolov5.cpp @@ -11,13 +11,6 @@ #define NMS_THRESH 0.4 #define CONF_THRESH 0.5 #define BATCH_SIZE 1 -#define MODEL "../yolov5.wts" - -#define NET s // s m l x -#define NETSTRUCT(str) createEngine_##str -#define CREATENET(net) NETSTRUCT(net) -#define STR1(x) #x -#define STR2(x) STR1(x) // stuff we know about the network and the input/output blobs static const int INPUT_H = Yolo::INPUT_H; @@ -27,10 +20,8 @@ static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detect const char* INPUT_BLOB_NAME = "data"; const char* OUTPUT_BLOB_NAME = "prob"; static Logger gLogger; -static float gw; -static float gd; -int get_width(int x, int divisor = 8) { +static int get_width(int x, float gw, int divisor = 8) { //return math.ceil(x / divisor) * divisor if (int(x * gw) % divisor == 0) { return int(x * gw); @@ -38,76 +29,75 @@ int get_width(int x, int divisor = 8) { return (int(x * gw / divisor) + 1) * divisor; } -int get_depth(int x) { +static int get_depth(int x, float gd) { if (x == 1) { return 1; - } - else { + } else { return round(x * gd) > 1 ? round(x * gd) : 1; } } -ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { +ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) { INetworkDefinition* network = builder->createNetworkV2(0U); // 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 }); assert(data); - std::map weightMap = loadWeights(MODEL); + std::map weightMap = loadWeights(wts_name); Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; /* ------ yolov5 backbone------ */ - auto focus0 = focus(network, weightMap, *data, 3, get_width(64), 3, "model.0"); - auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128), 3, 2, 1, "model.1"); - auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128), get_width(128), get_depth(3), true, 1, 0.5, "model.2"); - auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256), 3, 2, 1, "model.3"); - auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256), get_width(256), get_depth(9), true, 1, 0.5, "model.4"); - auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512), 3, 2, 1, "model.5"); - auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512), get_width(512), get_depth(9), true, 1, 0.5, "model.6"); - auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024), 3, 2, 1, "model.7"); - auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), get_width(1024), get_width(1024), 5, 9, 13, "model.8"); + auto focus0 = focus(network, weightMap, *data, 3, get_width(64, gw), 3, "model.0"); + auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1"); + auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2"); + auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); + auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(9, gd), true, 1, 0.5, "model.4"); + auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); + auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6"); + auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7"); + auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, 9, 13, "model.8"); /* ------ yolov5 head ------ */ - auto bottleneck_csp9 = C3(network, weightMap, *spp8->getOutput(0), get_width(1024), get_width(1024), get_depth(3), false, 1, 0.5, "model.9"); - auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), get_width(512), 1, 1, 1, "model.10"); + auto bottleneck_csp9 = C3(network, weightMap, *spp8->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.9"); + auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), get_width(512, gw), 1, 1, 1, "model.10"); - float* deval = reinterpret_cast(malloc(sizeof(float) * get_width(512) * 2 * 2)); - for (int i = 0; i < get_width(512) * 2 * 2; i++) { + float* deval = reinterpret_cast(malloc(sizeof(float) * get_width(512, gw) * 2 * 2)); + for (int i = 0; i < get_width(512, gw) * 2 * 2; i++) { deval[i] = 1.0; } - Weights deconvwts11{ DataType::kFLOAT, deval, get_width(512) * 2 * 2 }; - IDeconvolutionLayer* deconv11 = network->addDeconvolutionNd(*conv10->getOutput(0), get_width(512), DimsHW{ 2, 2 }, deconvwts11, emptywts); + Weights deconvwts11{ DataType::kFLOAT, deval, get_width(512, gw) * 2 * 2 }; + IDeconvolutionLayer* deconv11 = network->addDeconvolutionNd(*conv10->getOutput(0), get_width(512, gw), DimsHW{ 2, 2 }, deconvwts11, emptywts); deconv11->setStrideNd(DimsHW{ 2, 2 }); - deconv11->setNbGroups(get_width(512)); + deconv11->setNbGroups(get_width(512, gw)); weightMap["deconv11"] = deconvwts11; ITensor* inputTensors12[] = { deconv11->getOutput(0), bottleneck_csp6->getOutput(0) }; auto cat12 = network->addConcatenation(inputTensors12, 2); - auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), get_width(1024), get_width(512), get_depth(3), false, 1, 0.5, "model.13"); - auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256), 1, 1, 1, "model.14"); + auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.13"); + auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256, gw), 1, 1, 1, "model.14"); - Weights deconvwts15{ DataType::kFLOAT, deval, get_width(256) * 2 * 2 }; - IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), get_width(256), DimsHW{ 2, 2 }, deconvwts15, emptywts); + Weights deconvwts15{ DataType::kFLOAT, deval, get_width(256, gw) * 2 * 2 }; + IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), get_width(256, gw), DimsHW{ 2, 2 }, deconvwts15, emptywts); deconv15->setStrideNd(DimsHW{ 2, 2 }); - deconv15->setNbGroups(get_width(256)); + deconv15->setNbGroups(get_width(256, gw)); ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) }; auto cat16 = network->addConcatenation(inputTensors16, 2); - auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512), get_width(256), get_depth(3), false, 1, 0.5, "model.17"); + auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.17"); // yolo layer 0 IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256), 3, 2, 1, "model.18"); + auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 3, 2, 1, "model.18"); ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512), get_width(512), get_depth(3), false, 1, 0.5, "model.20"); + auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.20"); //yolo layer 1 IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512), 3, 2, 1, "model.21"); + auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512, gw), 3, 2, 1, "model.21"); ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024), get_width(1024), get_depth(3), false, 1, 0.5, "model.23"); + auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.23"); IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); @@ -143,403 +133,13 @@ ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilder return engine; } -// Creat the engine using only the API and not any parser. -ICudaEngine* createEngine_s(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { - INetworkDefinition* network = builder->createNetworkV2(0U); - - // 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 }); - assert(data); - - std::map weightMap = loadWeights("../yolov5s.wts"); - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - - // yolov5 backbone - auto focus0 = focus(network, weightMap, *data, 3, 32, 3, "model.0"); - auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), 64, 3, 2, 1, "model.1"); - auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), 64, 64, 1, true, 1, 0.5, "model.2"); - auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), 128, 3, 2, 1, "model.3"); - auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), 128, 128, 3, true, 1, 0.5, "model.4"); - auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), 256, 3, 2, 1, "model.5"); - auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), 256, 256, 3, true, 1, 0.5, "model.6"); - auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), 512, 3, 2, 1, "model.7"); - auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), 512, 512, 5, 9, 13, "model.8"); - - // yolov5 head - auto bottleneck_csp9 = C3(network, weightMap, *spp8->getOutput(0), 512, 512, 1, false, 1, 0.5, "model.9"); - auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), 256, 1, 1, 1, "model.10"); - - float *deval = reinterpret_cast(malloc(sizeof(float) * 256 * 2 * 2)); - for (int i = 0; i < 256 * 2 * 2; i++) { - deval[i] = 1.0; - } - Weights deconvwts11{ DataType::kFLOAT, deval, 256 * 2 * 2 }; - IDeconvolutionLayer* deconv11 = network->addDeconvolutionNd(*conv10->getOutput(0), 256, DimsHW{ 2, 2 }, deconvwts11, emptywts); - deconv11->setStrideNd(DimsHW{ 2, 2 }); - deconv11->setNbGroups(256); - weightMap["deconv11"] = deconvwts11; - - ITensor* inputTensors12[] = { deconv11->getOutput(0), bottleneck_csp6->getOutput(0) }; - auto cat12 = network->addConcatenation(inputTensors12, 2); - auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), 512, 256, 1, false, 1, 0.5, "model.13"); - auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), 128, 1, 1, 1, "model.14"); - - Weights deconvwts15{ DataType::kFLOAT, deval, 128 * 2 * 2 }; - IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), 128, DimsHW{ 2, 2 }, deconvwts15, emptywts); - deconv15->setStrideNd(DimsHW{ 2, 2 }); - deconv15->setNbGroups(128); - - ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) }; - auto cat16 = network->addConcatenation(inputTensors16, 2); - auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), 256, 128, 1, false, 1, 0.5, "model.17"); - IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), 128, 3, 2, 1, "model.18"); - ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; - auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), 256, 256, 1, false, 1, 0.5, "model.20"); - IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), 256, 3, 2, 1, "model.21"); - ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; - auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), 512, 512, 1, false, 1, 0.5, "model.23"); - IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - - auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); - yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); - network->markOutput(*yolo->getOutput(0)); - - // Build engine - builder->setMaxBatchSize(maxBatchSize); - config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB -#if defined(USE_FP16) - config->setFlag(BuilderFlag::kFP16); -#elif defined(USE_INT8) - std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; - assert(builder->platformHasFastInt8()); - config->setFlag(BuilderFlag::kINT8); - Int8EntropyCalibrator2 *calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME); - config->setInt8Calibrator(calibrator); -#endif - - std::cout << "Building engine, please wait for a while..." << std::endl; - ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); - std::cout << "Build engine successfully!" << std::endl; - - // Don't need the network any more - network->destroy(); - - // Release host memory - for (auto& mem : weightMap) - { - free((void*)(mem.second.values)); - - } - return engine; -} - -ICudaEngine* createEngine_m(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { - INetworkDefinition* network = builder->createNetworkV2(0U); - - // 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 }); - assert(data); - - std::map weightMap = loadWeights("../yolov5m.wts"); - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - - /* ------ yolov5 backbone------ */ - auto focus0 = focus(network, weightMap, *data, 3, 48, 3, "model.0"); - auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), 96, 3, 2, 1, "model.1"); - auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), 96, 96, 2, true, 1, 0.5, "model.2"); - auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), 192, 3, 2, 1, "model.3"); - auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), 192, 192, 6, true, 1, 0.5, "model.4"); - auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), 384, 3, 2, 1, "model.5"); - auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), 384, 384, 6, true, 1, 0.5, "model.6"); - auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), 768, 3, 2, 1, "model.7"); - auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), 768, 768, 5, 9, 13, "model.8"); - /* ------ yolov5 head ------ */ - auto bottleneck_csp9 = C3(network, weightMap, *spp8->getOutput(0), 768, 768, 2, false, 1, 0.5, "model.9"); - auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), 384, 1, 1, 1, "model.10"); - - float *deval = reinterpret_cast(malloc(sizeof(float) * 384 * 2 * 2)); - for (int i = 0; i < 384 * 2 * 2; i++) { - deval[i] = 1.0; - } - Weights deconvwts11{ DataType::kFLOAT, deval, 384 * 2 * 2 }; - IDeconvolutionLayer* deconv11 = network->addDeconvolutionNd(*conv10->getOutput(0), 384, DimsHW{ 2, 2 }, deconvwts11, emptywts); - deconv11->setStrideNd(DimsHW{ 2, 2 }); - deconv11->setNbGroups(384); - weightMap["deconv11"] = deconvwts11; - ITensor* inputTensors12[] = { deconv11->getOutput(0), bottleneck_csp6->getOutput(0) }; - auto cat12 = network->addConcatenation(inputTensors12, 2); - - auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), 768, 384, 2, false, 1, 0.5, "model.13"); - - auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), 192, 1, 1, 1, "model.14"); - - Weights deconvwts15{ DataType::kFLOAT, deval, 192 * 2 * 2 }; - IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), 192, DimsHW{ 2, 2 }, deconvwts15, emptywts); - deconv15->setStrideNd(DimsHW{ 2, 2 }); - deconv15->setNbGroups(192); - - ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) }; - auto cat16 = network->addConcatenation(inputTensors16, 2); - auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), 384, 192, 2, false, 1, 0.5, "model.17"); - - //yolo layer 0 - IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), 192, 3, 2, 1, "model.18"); - ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; - auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), 384, 384, 2, false, 1, 0.5, "model.20"); - - //yolo layer 1 - IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), 384, 3, 2, 1, "model.21"); - ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; - auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), 768, 768, 2, false, 1, 0.5, "model.23"); - // yolo layer 2 - IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - - auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); - yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); - network->markOutput(*yolo->getOutput(0)); - - // Build engine - builder->setMaxBatchSize(maxBatchSize); - config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB -#if defined(USE_FP16) - config->setFlag(BuilderFlag::kFP16); -#elif defined(USE_INT8) - std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; - assert(builder->platformHasFastInt8()); - config->setFlag(BuilderFlag::kINT8); - Int8EntropyCalibrator2 *calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME); - config->setInt8Calibrator(calibrator); -#endif - - std::cout << "Building engine, please wait for a while..." << std::endl; - ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); - std::cout << "Build engine successfully!" << std::endl; - - // Don't need the network any more - network->destroy(); - - // Release host memory - for (auto& mem : weightMap) - { - free((void*)(mem.second.values)); - } - - return engine; -} - -ICudaEngine* createEngine_l(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { - INetworkDefinition* network = builder->createNetworkV2(0U); - - // 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 }); - assert(data); - - std::map weightMap = loadWeights("../yolov5l.wts"); - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - - /* ------ yolov5 backbone------ */ - auto focus0 = focus(network, weightMap, *data, 3, 64, 3, "model.0"); - auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), 128, 3, 2, 1, "model.1"); - auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), 128, 128, 3, true, 1, 0.5, "model.2"); - auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), 256, 3, 2, 1, "model.3"); - auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), 256, 256, 9, true, 1, 0.5, "model.4"); - auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), 512, 3, 2, 1, "model.5"); - auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), 512, 512, 9, true, 1, 0.5, "model.6"); - auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), 1024, 3, 2, 1, "model.7"); - auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), 1024, 1024, 5, 9, 13, "model.8"); - - /* ------ yolov5 head ------ */ - auto bottleneck_csp9 = C3(network, weightMap, *spp8->getOutput(0), 1024, 1024, 3, false, 1, 0.5, "model.9"); - auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), 512, 1, 1, 1, "model.10"); - - float *deval = reinterpret_cast(malloc(sizeof(float) * 512 * 2 * 2)); - for (int i = 0; i < 512 * 2 * 2; i++) { - deval[i] = 1.0; - } - Weights deconvwts11{ DataType::kFLOAT, deval, 512 * 2 * 2 }; - IDeconvolutionLayer* deconv11 = network->addDeconvolutionNd(*conv10->getOutput(0), 512, DimsHW{ 2, 2 }, deconvwts11, emptywts); - deconv11->setStrideNd(DimsHW{ 2, 2 }); - deconv11->setNbGroups(512); - weightMap["deconv11"] = deconvwts11; - - ITensor* inputTensors12[] = { deconv11->getOutput(0), bottleneck_csp6->getOutput(0) }; - auto cat12 = network->addConcatenation(inputTensors12, 2); - auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), 1024, 512, 3, false, 1, 0.5, "model.13"); - auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), 256, 1, 1, 1, "model.14"); - - Weights deconvwts15{ DataType::kFLOAT, deval, 256 * 2 * 2 }; - IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), 256, DimsHW{ 2, 2 }, deconvwts15, emptywts); - deconv15->setStrideNd(DimsHW{ 2, 2 }); - deconv15->setNbGroups(256); - ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) }; - auto cat16 = network->addConcatenation(inputTensors16, 2); - - auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), 512, 256, 3, false, 1, 0.5, "model.17"); - - // yolo layer 0 - IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), 256, 3, 2, 1, "model.18"); - ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; - auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), 512, 512, 3, false, 1, 0.5, "model.20"); - //yolo layer 1 - IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), 512, 3, 2, 1, "model.21"); - ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; - auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), 1024, 1024, 3, false, 1, 0.5, "model.23"); - IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - - auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); - yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); - network->markOutput(*yolo->getOutput(0)); - - // Build engine - builder->setMaxBatchSize(maxBatchSize); - config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB -#if defined(USE_FP16) - config->setFlag(BuilderFlag::kFP16); -#elif defined(USE_INT8) - std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; - assert(builder->platformHasFastInt8()); - config->setFlag(BuilderFlag::kINT8); - Int8EntropyCalibrator2 *calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME); - config->setInt8Calibrator(calibrator); -#endif - - std::cout << "Building engine, please wait for a while..." << std::endl; - ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); - std::cout << "Build engine successfully!" << std::endl; - - // Don't need the network any more - network->destroy(); - - // Release host memory - for (auto& mem : weightMap) - { - free((void*)(mem.second.values)); - } - - return engine; -} - -ICudaEngine* createEngine_x(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { - INetworkDefinition* network = builder->createNetworkV2(0U); - - // 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 }); - assert(data); - - std::map weightMap = loadWeights("../yolov5x.wts"); - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - - /* ------ yolov5 backbone------ */ - auto focus0 = focus(network, weightMap, *data, 3, 80, 3, "model.0"); - auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), 160, 3, 2, 1, "model.1"); - auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), 160, 160, 4, true, 1, 0.5, "model.2"); - auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), 320, 3, 2, 1, "model.3"); - auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), 320, 320, 12, true, 1, 0.5, "model.4"); - auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), 640, 3, 2, 1, "model.5"); - auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), 640, 640, 12, true, 1, 0.5, "model.6"); - auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), 1280, 3, 2, 1, "model.7"); - auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), 1280, 1280, 5, 9, 13, "model.8"); - - /* ------- yolov5 head ------- */ - auto bottleneck_csp9 = C3(network, weightMap, *spp8->getOutput(0), 1280, 1280, 4, false, 1, 0.5, "model.9"); - auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), 640, 1, 1, 1, "model.10"); - - float *deval = reinterpret_cast(malloc(sizeof(float) * 640 * 2 * 2)); - for (int i = 0; i < 640 * 2 * 2; i++) { - deval[i] = 1.0; - } - Weights deconvwts11{ DataType::kFLOAT, deval, 640 * 2 * 2 }; - IDeconvolutionLayer* deconv11 = network->addDeconvolutionNd(*conv10->getOutput(0), 640, DimsHW{ 2, 2 }, deconvwts11, emptywts); - deconv11->setStrideNd(DimsHW{ 2, 2 }); - deconv11->setNbGroups(640); - weightMap["deconv11"] = deconvwts11; - - ITensor* inputTensors12[] = { deconv11->getOutput(0), bottleneck_csp6->getOutput(0) }; - auto cat12 = network->addConcatenation(inputTensors12, 2); - - auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), 1280, 640, 4, false, 1, 0.5, "model.13"); - auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), 320, 1, 1, 1, "model.14"); - - Weights deconvwts15{ DataType::kFLOAT, deval, 320 * 2 * 2 }; - IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), 320, DimsHW{ 2, 2 }, deconvwts15, emptywts); - deconv15->setStrideNd(DimsHW{ 2, 2 }); - deconv15->setNbGroups(320); - ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) }; - auto cat16 = network->addConcatenation(inputTensors16, 2); - - auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), 640, 320, 4, false, 1, 0.5, "model.17"); - - // yolo layer 0 - IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); - auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), 320, 3, 2, 1, "model.18"); - ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; - auto cat19 = network->addConcatenation(inputTensors19, 2); - auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), 640, 640, 4, false, 1, 0.5, "model.20"); - // yolo layer 1 - IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); - auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), 640, 3, 2, 1, "model.21"); - ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; - auto cat22 = network->addConcatenation(inputTensors22, 2); - auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), 1280, 1280, 4, false, 1, 0.5, "model.23"); - // yolo layer 2 - IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - - auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); - yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); - network->markOutput(*yolo->getOutput(0)); - - // Build engine - builder->setMaxBatchSize(maxBatchSize); - config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB -#if defined(USE_FP16) - config->setFlag(BuilderFlag::kFP16); -#elif defined(USE_INT8) - std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; - assert(builder->platformHasFastInt8()); - config->setFlag(BuilderFlag::kINT8); - Int8EntropyCalibrator2 *calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME); - config->setInt8Calibrator(calibrator); -#endif - - std::cout << "Building engine, please wait for a while..." << std::endl; - ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); - std::cout << "Build engine successfully!" << std::endl; - - // Don't need the network any more - network->destroy(); - - // Release host memory - for (auto& mem : weightMap) - { - free((void*)(mem.second.values)); - } - - return engine; -} - -void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) { +void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, float& gd, float& gw, std::string& wts_name) { // Create builder IBuilder* builder = createInferBuilder(gLogger); IBuilderConfig* config = builder->createBuilderConfig(); // Create model to populate the network, then set the outputs and create an engine - //ICudaEngine* engine = (CREATENET(NET))(maxBatchSize, builder, config, DataType::kFLOAT); - ICudaEngine* engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT); - //ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT); + ICudaEngine* engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name); assert(engine != nullptr); // Serialize the engine @@ -548,6 +148,7 @@ void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) { // Close everything down engine->destroy(); builder->destroy(); + config->destroy(); } void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffers, float* input, float* output, int batchSize) { @@ -558,19 +159,57 @@ void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffer cudaStreamSynchronize(stream); } +bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, float& gd, float& gw, std::string& img_dir) { + if (argc < 4) return false; + if (std::string(argv[1]) == "-s" && (argc == 5 || argc == 7)) { + wts = std::string(argv[2]); + engine = std::string(argv[3]); + auto net = std::string(argv[4]); + if (net == "s") { + gd = 0.33; + gw = 0.50; + } else if (net == "m") { + gd = 0.67; + gw = 0.75; + } else if (net == "l") { + gd = 1.0; + gw = 1.0; + } else if (net == "x") { + gd = 1.33; + gw = 1.25; + } else if (net == "c" && argc == 7) { + gd = atof(argv[5]); + gw = atof(argv[6]); + } else { + return false; + } + } else if (std::string(argv[1]) == "-d" && argc == 4) { + engine = std::string(argv[2]); + img_dir = std::string(argv[3]); + } else { + return false; + } + return true; +} + int main(int argc, char** argv) { cudaSetDevice(DEVICE); + + std::string wts_name = ""; + std::string engine_name = ""; + float gd = 0.0f, gw = 0.0f; + std::string img_dir; + if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) { + std::cerr << "arguments not right!" << std::endl; + std::cerr << "./yolov5 -s [.wts] [.engine] [s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; + std::cerr << "./yolov5 -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + return -1; + } + // create a model using the API directly and serialize it to a stream - char *trtModelStream{ nullptr }; - size_t size{ 0 }; - std::string engine_name = STR2(NET); - engine_name = "yolov5" + engine_name + ".engine"; - if (argc == 4 && std::string(argv[1]) == "-s") { - //The third arg is depth and the fourth is width - gw = atof(argv[3]); - gd = atof(argv[2]); + if (!wts_name.empty()) { IHostMemory* modelStream{ nullptr }; - APIToModel(BATCH_SIZE, &modelStream); + APIToModel(BATCH_SIZE, &modelStream, gd, gw, wts_name); assert(modelStream != nullptr); std::ofstream p(engine_name, std::ios::binary); if (!p) { @@ -580,27 +219,27 @@ int main(int argc, char** argv) { p.write(reinterpret_cast(modelStream->data()), modelStream->size()); modelStream->destroy(); return 0; - } else if (argc == 3 && std::string(argv[1]) == "-d") { - std::ifstream file(engine_name, std::ios::binary); - if (file.good()) { - file.seekg(0, file.end); - size = file.tellg(); - file.seekg(0, file.beg); - trtModelStream = new char[size]; - assert(trtModelStream); - file.read(trtModelStream, size); - file.close(); - } - } else { - std::cerr << "arguments not right!" << std::endl; - std::cerr << "./yolov5 -s // serialize model to plan file" << std::endl; - std::cerr << "./yolov5 -d ../samples // deserialize plan file and run inference" << std::endl; - return -1; } + // deserialize the .engine and run inference + std::ifstream file(engine_name, std::ios::binary); + if (!file.good()) { + std::cerr << "read " << engine_name << " error!" << std::endl; + return -1; + } + char *trtModelStream = nullptr; + size_t size = 0; + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + trtModelStream = new char[size]; + assert(trtModelStream); + file.read(trtModelStream, size); + file.close(); + std::vector file_names; - if (read_files_in_dir(argv[2], file_names) < 0) { - std::cout << "read_files_in_dir failed." << std::endl; + if (read_files_in_dir(img_dir.c_str(), file_names) < 0) { + std::cerr << "read_files_in_dir failed." << std::endl; return -1; } @@ -636,7 +275,7 @@ int main(int argc, char** argv) { fcount++; 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]); + cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]); if (img.empty()) continue; cv::Mat pr_img = preprocess_img(img, INPUT_W, INPUT_H); // letterbox BGR to RGB int i = 0; @@ -665,7 +304,7 @@ int main(int argc, char** argv) { 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]); + cv::Mat img = cv::imread(img_dir + "/" + 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);