diff --git a/yolov5/README.md b/yolov5/README.md index 836031d..d1272b9 100644 --- a/yolov5/README.md +++ b/yolov5/README.md @@ -4,9 +4,10 @@ The Pytorch implementation is [ultralytics/yolov5](https://github.com/ultralytic ## Different versions of yolov5 -Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1 and v4.0. +Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1, v4.0 and v5.0. -- 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`, then follow how-to-run in current page. +- For yolov5 v5.0, download .pt from [yolov5 release v5.0](https://github.com/ultralytics/yolov5/releases/tag/v5.0), `git clone -b v5.0 https://github.com/ultralytics/yolov5.git` and `git clone https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in current page. +- 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 -b yolov5-v4.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v4.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v4.0/yolov5). - 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 -b yolov5-v3.1 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v3.1](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v3.1/yolov5). - 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 -b yolov5-v3.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v3.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v3.0/yolov5). - 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 -b yolov5-v2.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v2.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v2.0/yolov5). @@ -28,20 +29,20 @@ Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1 and v4.0. 1. generate .wts from pytorch with .pt, or download .wts from model zoo ``` -// 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 -python gen_wts.py +git clone -b v5.0 https://github.com/ultralytics/yolov5.git +git clone https://github.com/wang-xinyu/tensorrtx.git +// download https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt +cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5 +cd {ultralytics}/yolov5 +python gen_wts.py yolov5s.pt // a file 'yolov5s.wts' will be generated. ``` 2. build tensorrtx/yolov5 and run ``` -// put yolov5s.wts into tensorrtx/yolov5 -// go to tensorrtx/yolov5 +cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/ +cd {tensorrtx}/yolov5/ // update CLASS_NUM in yololayer.h if your model is trained on custom dataset mkdir build cd build diff --git a/yolov5/gen_wts.py b/yolov5/gen_wts.py index dbb2dde..501fd63 100644 --- a/yolov5/gen_wts.py +++ b/yolov5/gen_wts.py @@ -1,14 +1,16 @@ import torch import struct +import sys from utils.torch_utils import select_device # Initialize device = select_device('cpu') +pt_file = sys.argv[1] # Load model -model = torch.load('weights/yolov5s.pt', map_location=device)['model'].float() # load to FP32 +model = torch.load(pt_file, map_location=device)['model'].float() # load to FP32 model.to(device).eval() -with open('yolov5s.wts', 'w') as f: +with open(pt_file.split('.')[0] + '.wts', 'w') as f: f.write('{}\n'.format(len(model.state_dict().keys()))) for k, v in model.state_dict().items(): vr = v.reshape(-1).cpu().numpy() diff --git a/yolov5/yolov5.cpp b/yolov5/yolov5.cpp index 66cc26b..2fcf50b 100644 --- a/yolov5/yolov5.cpp +++ b/yolov5/yolov5.cpp @@ -128,6 +128,111 @@ ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilder return engine; } +ICudaEngine* build_engine_p6(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(wts_name); + + /* ------ yolov5 backbone------ */ + 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 c3_2 = 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, *c3_2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3"); + auto c3_4 = 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, *c3_4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5"); + auto c3_6 = 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, *c3_6->getOutput(0), get_width(768, gw), 3, 2, 1, "model.7"); + auto c3_8 = C3(network, weightMap, *conv7->getOutput(0), get_width(768, gw), get_width(768, gw), get_depth(3, gd), true, 1, 0.5, "model.8"); + auto conv9 = convBlock(network, weightMap, *c3_8->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.9"); + auto spp10 = SPP(network, weightMap, *conv9->getOutput(0), get_width(1024, gw), get_width(1024, gw), 3, 5, 7, "model.10"); + auto c3_11 = C3(network, weightMap, *spp10->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.11"); + + /* ------ yolov5 head ------ */ + auto conv12 = convBlock(network, weightMap, *c3_11->getOutput(0), get_width(768, gw), 1, 1, 1, "model.12"); + auto upsample13 = network->addResize(*conv12->getOutput(0)); + assert(upsample13); + upsample13->setResizeMode(ResizeMode::kNEAREST); + upsample13->setOutputDimensions(c3_8->getOutput(0)->getDimensions()); + ITensor* inputTensors14[] = { upsample13->getOutput(0), c3_8->getOutput(0) }; + auto cat14 = network->addConcatenation(inputTensors14, 2); + auto c3_15 = C3(network, weightMap, *cat14->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.15"); + + auto conv16 = convBlock(network, weightMap, *c3_15->getOutput(0), get_width(512, gw), 1, 1, 1, "model.16"); + auto upsample17 = network->addResize(*conv16->getOutput(0)); + assert(upsample17); + upsample17->setResizeMode(ResizeMode::kNEAREST); + upsample17->setOutputDimensions(c3_6->getOutput(0)->getDimensions()); + ITensor* inputTensors18[] = { upsample17->getOutput(0), c3_6->getOutput(0) }; + auto cat18 = network->addConcatenation(inputTensors18, 2); + auto c3_19 = C3(network, weightMap, *cat18->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.19"); + + auto conv20 = convBlock(network, weightMap, *c3_19->getOutput(0), get_width(256, gw), 1, 1, 1, "model.20"); + auto upsample21 = network->addResize(*conv20->getOutput(0)); + assert(upsample21); + upsample21->setResizeMode(ResizeMode::kNEAREST); + upsample21->setOutputDimensions(c3_4->getOutput(0)->getDimensions()); + ITensor* inputTensors21[] = { upsample21->getOutput(0), c3_4->getOutput(0) }; + auto cat22 = network->addConcatenation(inputTensors21, 2); + auto c3_23 = C3(network, weightMap, *cat22->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.23"); + + auto conv24 = convBlock(network, weightMap, *c3_23->getOutput(0), get_width(256, gw), 3, 2, 1, "model.24"); + ITensor* inputTensors25[] = { conv24->getOutput(0), conv20->getOutput(0) }; + auto cat25 = network->addConcatenation(inputTensors25, 2); + auto c3_26 = C3(network, weightMap, *cat25->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.26"); + + auto conv27 = convBlock(network, weightMap, *c3_26->getOutput(0), get_width(512, gw), 3, 2, 1, "model.27"); + ITensor* inputTensors28[] = { conv27->getOutput(0), conv16->getOutput(0) }; + auto cat28 = network->addConcatenation(inputTensors28, 2); + auto c3_29 = C3(network, weightMap, *cat28->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.29"); + + auto conv30 = convBlock(network, weightMap, *c3_29->getOutput(0), get_width(768, gw), 3, 2, 1, "model.30"); + ITensor* inputTensors31[] = { conv30->getOutput(0), conv12->getOutput(0) }; + auto cat31 = network->addConcatenation(inputTensors31, 2); + auto c3_32 = C3(network, weightMap, *cat31->getOutput(0), get_width(2048, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.32"); + + /* ------ detect ------ */ + IConvolutionLayer* det0 = network->addConvolutionNd(*c3_23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.0.weight"], weightMap["model.33.m.0.bias"]); + IConvolutionLayer* det1 = network->addConvolutionNd(*c3_26->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.1.weight"], weightMap["model.33.m.1.bias"]); + IConvolutionLayer* det2 = network->addConvolutionNd(*c3_29->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.2.weight"], weightMap["model.33.m.2.bias"]); + IConvolutionLayer* det3 = network->addConvolutionNd(*c3_32->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.3.weight"], weightMap["model.33.m.3.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, float& gd, float& gw, std::string& wts_name) { // Create builder IBuilder* builder = createInferBuilder(gLogger);