From e065c7508dd306dbd5d2f86d5d26b2855146ea74 Mon Sep 17 00:00:00 2001 From: WuxinrongY <53141838+WuxinrongY@users.noreply.github.com> Date: Fri, 21 Jun 2024 15:44:24 +0800 Subject: [PATCH] yolov9 t s m (#1541) * yolov9 t s m * change chinese --- yolov9/README.md | 38 +- yolov9/demo.cpp | 45 +- yolov9/include/block.h | 4 + yolov9/include/model.h | 24 +- yolov9/src/block.cpp | 37 +- yolov9/src/model.cpp | 911 ++++++++++++++++++++++++++++++++--------- 6 files changed, 830 insertions(+), 229 deletions(-) diff --git a/yolov9/README.md b/yolov9/README.md index aebf19e..27250c2 100644 --- a/yolov9/README.md +++ b/yolov9/README.md @@ -7,22 +7,16 @@ The Pytorch implementation is [WongKinYiu/yolov9](https://github.com/WongKinYiu/ ## Progress -- [x] YOLOv9-c: - - [x] FP32 - - [x] FP16 - - [x] INT8 -- [x] YOLOv9-e: - - [x] FP32 - - [x] FP16 - - [x] INT8 -- [x] GELAN-c: - - [x] FP32 - - [x] FP16 - - [x] INT8 -- [x] GELAN-e: - - [x] FP32 - - [x] FP16 - - [x] INT8 +- [x] YOLOv9-t +- [x] YOLOv9-t-convert(gelan) +- [x] YOLOv9-s +- [x] YOLOv9-s-convert(gelan) +- [x] YOLOv9-m +- [x] YOLOv9-m-convert(gelan) +- [x] YOLOv9-c +- [x] YOLOv9-c-convert(gelan) +- [x] YOLOv9-e +- [x] YOLOv9-e-convert(gelan) ## Requirements @@ -35,8 +29,16 @@ The speed test is done on a desktop with R7-5700G CPU and RTX 4060Ti GPU. The in | frame | Model | FP32 | FP16 | INT8 | | --- | --- | --- | --- | --- | -| pytorch | YOLOv9-c | - | 15.5ms | - | -| pytorch | YOLOv9-e | - | 19.7ms | - | +| tensorrt | YOLOv5-n | -ms | 0.58ms | -ms | +| tensorrt | YOLOv5-s | -ms | 0.90ms | -ms | +| tensorrt | YOLOv5-m | -ms | 1.9ms | -ms | +| tensorrt | YOLOv5-l | -ms | 2.8ms | -ms | +| tensorrt | YOLOv5-x | -ms | 5.1ms | -ms | +| tensorrt | YOLOv9-t-convert | -ms | 1.37ms | -ms | +| tensorrt | YOLOv9-s | -ms | 1.78ms | -ms | +| tensorrt | YOLOv9-s-convert | -ms | 1.78ms | -ms | +| tensorrt | YOLOv9-m | -ms | 3.1ms | -ms | +| tensorrt | YOLOv9-m-convert | -ms | 2.8ms | -ms | | tensorrt | YOLOv9-c | 13.5ms | 4.6ms | 3.0ms | | tensorrt | YOLOv9-e | 8.3ms | 3.2ms | 2.15ms | diff --git a/yolov9/demo.cpp b/yolov9/demo.cpp index d193f8c..6c4ab75 100644 --- a/yolov9/demo.cpp +++ b/yolov9/demo.cpp @@ -19,17 +19,32 @@ void serialize_engine(unsigned int max_batchsize, std::string& wts_name, std::st // Create model to populate the network, then set the outputs and create an engine IHostMemory* serialized_engine = nullptr; - if (sub_type == "e") { - serialized_engine = build_engine_yolov9_e(max_batchsize, builder, config, DataType::kFLOAT, wts_name); + if (sub_type == "t") { + serialized_engine = build_engine_yolov9_t(max_batchsize, builder, config, DataType::kFLOAT, wts_name, false); + } else if (sub_type == "s") { + serialized_engine = build_engine_yolov9_s(max_batchsize, builder, config, DataType::kFLOAT, wts_name, false); + } else if (sub_type == "m") { + serialized_engine = build_engine_yolov9_m(max_batchsize, builder, config, DataType::kFLOAT, wts_name, false); } else if (sub_type == "c") { serialized_engine = build_engine_yolov9_c(max_batchsize, builder, config, DataType::kFLOAT, wts_name); - } else if (sub_type == "ge") { - serialized_engine = build_engine_gelan_e(max_batchsize, builder, config, DataType::kFLOAT, wts_name); + } else if (sub_type == "e") { + serialized_engine = build_engine_yolov9_e(max_batchsize, builder, config, DataType::kFLOAT, wts_name); + } + + else if (sub_type == "gt") { + serialized_engine = build_engine_yolov9_t(max_batchsize, builder, config, DataType::kFLOAT, wts_name, true); + } else if (sub_type == "gs") { + serialized_engine = build_engine_yolov9_s(max_batchsize, builder, config, DataType::kFLOAT, wts_name, true); + } else if (sub_type == "gm") { + serialized_engine = build_engine_yolov9_m(max_batchsize, builder, config, DataType::kFLOAT, wts_name, true); } else if (sub_type == "gc") { serialized_engine = build_engine_gelan_c(max_batchsize, builder, config, DataType::kFLOAT, wts_name); + } else if (sub_type == "ge") { + serialized_engine = build_engine_gelan_e(max_batchsize, builder, config, DataType::kFLOAT, wts_name); } else { return; } + assert(serialized_engine != nullptr); std::ofstream p(engine_name, std::ios::binary); @@ -114,19 +129,19 @@ int main(int argc, char** argv) { cudaSetDevice(kGpuId); std::string wts_name = ""; - std::string engine_name = ""; - std::string img_dir = ""; - std::string sub_type = ""; + std::string engine_name = "../yolov9-m-converted.engine"; + std::string img_dir = "../images"; + std::string sub_type = "m"; // speed test or inference - // const int speed_test_iter = 1000; - const int speed_test_iter = 1; + const int speed_test_iter = 1000; + // const int speed_test_iter = 1; - if (!parse_args(argc, argv, wts_name, engine_name, img_dir, sub_type)) { - std::cerr << "Arguments not right!" << std::endl; - std::cerr << "./yolov9 -s [.wts] [.engine] [c/e/gc/ge] // serialize model to plan file" << std::endl; - std::cerr << "./yolov9 -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; - return -1; - } + // if (!parse_args(argc, argv, wts_name, engine_name, img_dir, sub_type)) { + // std::cerr << "Arguments not right!" << std::endl; + // std::cerr << "./yolov9 -s [.wts] [.engine] [s/m/c/e/gt/gs/gm/gc/ge] // serialize model to plan file" << std::endl; + // std::cerr << "./yolov9 -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 file if (!wts_name.empty()) { diff --git a/yolov9/include/block.h b/yolov9/include/block.h index ca7f79f..f785776 100644 --- a/yolov9/include/block.h +++ b/yolov9/include/block.h @@ -22,10 +22,14 @@ std::vector> getAnchors(std::map& weigh // ---------------------------------------------------------------- nvinfer1::ILayer* convBnSiLU(nvinfer1::INetworkDefinition* network, std::map& weightMap, nvinfer1::ITensor& input, int ch, int k, int s, int p, std::string lname, int g = 1); +ILayer* ELAN1(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c1, int c2, + int c3, int c4, std::string lname); ILayer* RepNCSPELAN4(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c1, int c2, int c3, int c4, int c5, std::string lname); ILayer* ADown(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c2, std::string lname); +ILayer* AConv(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c2, + std::string lname); std::vector CBLinear(INetworkDefinition* network, std::map& weightMap, ITensor& input, std::vector c2s, int k, int s, int p, int g, std::string lname); ILayer* CBFuse(INetworkDefinition* network, std::vector> input, std::vector idx, diff --git a/yolov9/include/model.h b/yolov9/include/model.h index ced2694..66601be 100644 --- a/yolov9/include/model.h +++ b/yolov9/include/model.h @@ -2,16 +2,32 @@ #include #include -nvinfer1::IHostMemory* build_engine_yolov9_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, +// yolov9 +nvinfer1::IHostMemory* build_engine_yolov9_t(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, - std::string& wts_name); + std::string& wts_name, bool isConvert = false); +nvinfer1::IHostMemory* build_engine_yolov9_s(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, + nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, + std::string& wts_name, bool isConvert = false); +nvinfer1::IHostMemory* build_engine_yolov9_m(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, + nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, + std::string& wts_name, bool isConvert = false); nvinfer1::IHostMemory* build_engine_yolov9_c(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, std::string& wts_name); - -nvinfer1::IHostMemory* build_engine_gelan_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, +nvinfer1::IHostMemory* build_engine_yolov9_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, + nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, + std::string& wts_name); +// gelan +nvinfer1::IHostMemory* build_engine_gelan_t(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, + nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, + std::string& wts_name); +nvinfer1::IHostMemory* build_engine_gelan_m(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, std::string& wts_name); nvinfer1::IHostMemory* build_engine_gelan_c(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, std::string& wts_name); +nvinfer1::IHostMemory* build_engine_gelan_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, + nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, + std::string& wts_name); diff --git a/yolov9/src/block.cpp b/yolov9/src/block.cpp index c1f2393..d3f15ef 100644 --- a/yolov9/src/block.cpp +++ b/yolov9/src/block.cpp @@ -204,11 +204,34 @@ ILayer* RepNCSP(INetworkDefinition* network, std::map& wei auto cv3 = convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv3", 1); return cv3; } + +ILayer* ELAN1(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c1, int c2, + int c3, int c4, std::string lname) { + auto cv1 = convBnSiLU(network, weightMap, input, c3, 1, 1, 0, lname + ".cv1", 1); + // chunk(2, 1) + + nvinfer1::Dims d = cv1->getOutput(0)->getDimensions(); + nvinfer1::ISliceLayer* split1 = + network->addSlice(*cv1->getOutput(0), nvinfer1::Dims3{0, 0, 0}, nvinfer1::Dims3{d.d[0] / 2, d.d[1], d.d[2]}, + nvinfer1::Dims3{1, 1, 1}); + nvinfer1::ISliceLayer* split2 = + network->addSlice(*cv1->getOutput(0), nvinfer1::Dims3{d.d[0] / 2, 0, 0}, + nvinfer1::Dims3{d.d[0] / 2, d.d[1], d.d[2]}, nvinfer1::Dims3{1, 1, 1}); + auto cv2 = convBnSiLU(network, weightMap, *split2->getOutput(0), c4, 3, 1, 1, lname + ".cv2", 1); + + auto cv3 = convBnSiLU(network, weightMap, *cv2->getOutput(0), c4, 3, 1, 1, lname + ".cv3", 1); + + ITensor* inputTensors[] = {split1->getOutput(0), split2->getOutput(0), cv2->getOutput(0), cv3->getOutput(0)}; + auto cat = network->addConcatenation(inputTensors, 4); + auto cv4 = convBnSiLU(network, weightMap, *cat->getOutput(0), c2, 1, 1, 0, lname + ".cv4", 1); + return cv4; +} + ILayer* RepNCSPELAN4(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c1, int c2, int c3, int c4, int c5, std::string lname) { auto cv1 = convBnSiLU(network, weightMap, input, c3, 1, 1, 0, lname + ".cv1", 1); - // 将cv1的输出分成两部分 chunk(2, 1) + // chunk(2, 1) nvinfer1::Dims d = cv1->getOutput(0)->getDimensions(); nvinfer1::ISliceLayer* split1 = @@ -230,6 +253,14 @@ ILayer* RepNCSPELAN4(INetworkDefinition* network, std::map return cv4; } +ILayer* AConv(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c2, + std::string lname) { + auto pool = network->addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{2, 2}); + pool->setStrideNd(DimsHW{1, 1}); + pool->setPaddingNd(DimsHW{0, 0}); + auto cv1 = convBnSiLU(network, weightMap, *pool->getOutput(0), c2, 3, 2, 1, lname + ".cv1", 1); + return cv1; +} ILayer* ADown(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c2, std::string lname) { int c_ = c2 / 2; @@ -426,7 +457,9 @@ std::vector DualDDetect(INetworkDefinition* network, std:: std::vector DDetect(INetworkDefinition* network, std::map& weightMap, std::vector dets, int cls, std::vector ch, std::string lname) { int c2 = std::max(int(ch[0] / 4), int(16 * 4)); - int c3 = std::max(ch[0], std::min(cls * 2, 128)); + // max((ch[0], min((self.nc * 2, 128)))) + // int c3 = std::max(ch[0], std::min(cls * 2, 128)); + int c3 = std::max(ch[0], std::min(cls, 128)); int reg_max = 16; std::vector bboxlayers; diff --git a/yolov9/src/model.cpp b/yolov9/src/model.cpp index d450fb5..94e6389 100644 --- a/yolov9/src/model.cpp +++ b/yolov9/src/model.cpp @@ -22,6 +22,723 @@ void Calibrator(IBuilder* builder, IBuilderConfig* config) { } #endif +IHostMemory* build_engine_yolov9_t(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, + std::string& wts_name, bool isConvert) { + /* ------ Create the builder ------ */ + INetworkDefinition* network = builder->createNetworkV2(0U); + + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{3, kInputH, kInputW}); + assert(data); + std::map weightMap = loadWeights(wts_name); + + // # conv down + auto conv_1 = convBnSiLU(network, weightMap, *data, 16, 3, 2, 1, "model.0", 1); + // # conv down + auto conv_2 = convBnSiLU(network, weightMap, *conv_1->getOutput(0), 32, 3, 2, 1, "model.1"); + // # elan-1 block + auto repncspelan_3 = ELAN1(network, weightMap, *conv_2->getOutput(0), 32, 32, 32, 16, "model.2"); + // # avg-conv down + // [-1, 1, ADown, [256]], # 4-P3/8 + auto adown_4 = AConv(network, weightMap, *repncspelan_3->getOutput(0), 64, "model.3"); + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 5 + auto repncspelan_5 = RepNCSPELAN4(network, weightMap, *adown_4->getOutput(0), 64, 64, 64, 32, 3, "model.4"); + // # avg-conv down + // [-1, 1, ADown, [512]], # 6-P4/16 + auto adown_6 = AConv(network, weightMap, *repncspelan_5->getOutput(0), 96, "model.5"); + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 7 + auto repncspelan_7 = RepNCSPELAN4(network, weightMap, *adown_6->getOutput(0), 96, 96, 96, 48, 3, "model.6"); + // # avg-conv down + // [-1, 1, ADown, [512]], # 8-P5/32 + auto adown_8 = AConv(network, weightMap, *repncspelan_7->getOutput(0), 128, "model.7"); + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 9 + auto repncspelan_9 = RepNCSPELAN4(network, weightMap, *adown_8->getOutput(0), 128, 128, 128, 64, 3, "model.8"); + // # elan-spp block + // [-1, 1, SPPELAN, [512, 256]], # 10 + auto sppelan_10 = SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 128, 128, 64, "model.9"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_11 = network->addResize(*sppelan_10->getOutput(0)); + upsample_11->setResizeMode(ResizeMode::kNEAREST); + const float scales_11[] = {1.0, 2.0, 2.0}; + upsample_11->setScales(scales_11, 3); + // [[-1, 7], 1, Concat, [1]], # cat backbone P4 + ITensor* input_tensor_12[] = {upsample_11->getOutput(0), repncspelan_7->getOutput(0)}; + auto cat_12 = network->addConcatenation(input_tensor_12, 2); + + // # elan-2 block + auto repncspelan_13 = RepNCSPELAN4(network, weightMap, *cat_12->getOutput(0), 288, 96, 96, 48, 3, "model.12"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_14 = network->addResize(*repncspelan_13->getOutput(0)); + upsample_14->setResizeMode(ResizeMode::kNEAREST); + const float scales_14[] = {1.0, 2.0, 2.0}; + upsample_14->setScales(scales_14, 3); + // [[-1, 5], 1, Concat, [1]], # cat backbone P3 + ITensor* input_tensor_15[] = {upsample_14->getOutput(0), repncspelan_5->getOutput(0)}; + auto cat_15 = network->addConcatenation(input_tensor_15, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]], # 16 (P3/8-small) + auto repncspelan_16 = RepNCSPELAN4(network, weightMap, *cat_15->getOutput(0), 192, 64, 64, 32, 3, "model.15"); + + // # avg-conv-down merge + // [-1, 1, ADown, [256]], + auto adown_17 = AConv(network, weightMap, *repncspelan_16->getOutput(0), 48, "model.16"); + // [[-1, 13], 1, Concat, [1]], # cat head P4 + ITensor* input_tensor_18[] = {adown_17->getOutput(0), repncspelan_13->getOutput(0)}; + auto cat_18 = network->addConcatenation(input_tensor_18, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 19 (P4/16-medium) + auto repncspelan_19 = RepNCSPELAN4(network, weightMap, *cat_18->getOutput(0), 144, 96, 96, 48, 3, "model.18"); + + // # avg-conv-down merge + // [-1, 1, ADown, [512]], + auto adown_20 = AConv(network, weightMap, *repncspelan_19->getOutput(0), 64, "model.19"); + // [[-1, 10], 1, Concat, [1]], # cat head P5 + ITensor* input_tensor_21[] = {adown_20->getOutput(0), sppelan_10->getOutput(0)}; + auto cat_21 = network->addConcatenation(input_tensor_21, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 22 (P5/32-large) + auto repncspelan_22 = RepNCSPELAN4(network, weightMap, *cat_21->getOutput(0), 256, 128, 128, 64, 3, "model.21"); + + std::vector head; + if (!isConvert) { + // # elan-spp block + auto sppelan_23 = SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 512, 128, 64, "model.22"); + + // # up-concat merge + auto upsample_24 = network->addResize(*sppelan_23->getOutput(0)); + upsample_24->setResizeMode(ResizeMode::kNEAREST); + const float scales_24[] = {1.0, 2.0, 2.0}; + upsample_24->setScales(scales_24, 3); + // [[-1, 6], 1, Concat, [1]], # cat backbone P4 + ITensor* input_tensor_25[] = {upsample_24->getOutput(0), repncspelan_7->getOutput(0)}; + auto cat_25 = network->addConcatenation(input_tensor_25, 2); + + // # elan-2 block + auto repncspelan_26 = RepNCSPELAN4(network, weightMap, *cat_25->getOutput(0), 384, 96, 96, 48, 3, "model.25"); + + // # up-concat merge + auto upsample_27 = network->addResize(*repncspelan_26->getOutput(0)); + upsample_27->setResizeMode(ResizeMode::kNEAREST); + const float scales_27[] = {1.0, 2.0, 2.0}; + upsample_27->setScales(scales_27, 3); + // [[-1, 4], 1, Concat, [1]], # cat backbone P3 + ITensor* input_tensor_28[] = {upsample_27->getOutput(0), repncspelan_5->getOutput(0)}; + auto cat_28 = network->addConcatenation(input_tensor_28, 2); + + // # elan-2 block + auto repncspelan_29 = RepNCSPELAN4(network, weightMap, *cat_28->getOutput(0), 256, 64, 64, 32, 3, "model.28"); + head = DualDDetect(network, weightMap, std::vector{repncspelan_16, repncspelan_19, repncspelan_22}, + kNumClass, {64, 96, 128}, "model.29"); + } else { + head = DDetect(network, weightMap, std::vector{repncspelan_16, repncspelan_19, repncspelan_22}, + kNumClass, {64, 96, 128}, "model.22"); + } + + nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, head, false); + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + builder->setMaxBatchSize(kBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); + +#if defined(USE_FP16) + config->setFlag(nvinfer1::BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(nvinfer1::BuilderFlag::kINT8); + auto* calibrator = + new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + delete network; + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return serialized_model; +} + +IHostMemory* build_engine_yolov9_s(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, + std::string& wts_name, bool isConvert) { + /* ------ Create the builder ------ */ + INetworkDefinition* network = builder->createNetworkV2(0U); + + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{3, kInputH, kInputW}); + assert(data); + std::map weightMap = loadWeights(wts_name); + + // # conv down + auto conv_1 = convBnSiLU(network, weightMap, *data, 32, 3, 2, 1, "model.0", 1); + // # conv down + auto conv_2 = convBnSiLU(network, weightMap, *conv_1->getOutput(0), 64, 3, 2, 1, "model.1"); + // # elan-1 block + auto repncspelan_3 = ELAN1(network, weightMap, *conv_2->getOutput(0), 32, 64, 64, 32, "model.2"); + // # avg-conv down + auto adown_4 = AConv(network, weightMap, *repncspelan_3->getOutput(0), 128, "model.3"); + // # elan-2 block + auto repncspelan_5 = RepNCSPELAN4(network, weightMap, *adown_4->getOutput(0), 128, 128, 128, 64, 3, "model.4"); + // # avg-conv down + auto adown_6 = AConv(network, weightMap, *repncspelan_5->getOutput(0), 192, "model.5"); + // # elan-2 block + auto repncspelan_7 = RepNCSPELAN4(network, weightMap, *adown_6->getOutput(0), 192, 192, 192, 96, 3, "model.6"); + // # avg-conv down + auto adown_8 = AConv(network, weightMap, *repncspelan_7->getOutput(0), 256, "model.7"); + // # elan-2 block + auto repncspelan_9 = RepNCSPELAN4(network, weightMap, *adown_8->getOutput(0), 256, 256, 256, 128, 3, "model.8"); + // # elan-spp block + auto sppelan_10 = SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 512, 256, 128, "model.9"); + + // # up-concat merge + auto upsample_11 = network->addResize(*sppelan_10->getOutput(0)); + upsample_11->setResizeMode(ResizeMode::kNEAREST); + const float scales_11[] = {1.0, 2.0, 2.0}; + upsample_11->setScales(scales_11, 3); + // [[-1, 7], 1, Concat, [1]], # cat backbone P4 + ITensor* input_tensor_12[] = {upsample_11->getOutput(0), repncspelan_7->getOutput(0)}; + auto cat_12 = network->addConcatenation(input_tensor_12, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 13 + auto repncspelan_13 = RepNCSPELAN4(network, weightMap, *cat_12->getOutput(0), 192, 192, 192, 96, 3, "model.12"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_14 = network->addResize(*repncspelan_13->getOutput(0)); + upsample_14->setResizeMode(ResizeMode::kNEAREST); + const float scales_14[] = {1.0, 2.0, 2.0}; + upsample_14->setScales(scales_14, 3); + // [[-1, 5], 1, Concat, [1]], # cat backbone P3 + ITensor* input_tensor_15[] = {upsample_14->getOutput(0), repncspelan_5->getOutput(0)}; + auto cat_15 = network->addConcatenation(input_tensor_15, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]], # 16 (P3/8-small) + auto repncspelan_16 = RepNCSPELAN4(network, weightMap, *cat_15->getOutput(0), 128, 128, 128, 64, 3, "model.15"); + + // # avg-conv-down merge + // [-1, 1, ADown, [256]], + auto adown_17 = AConv(network, weightMap, *repncspelan_16->getOutput(0), 96, "model.16"); + // [[-1, 13], 1, Concat, [1]], # cat head P4 + ITensor* input_tensor_18[] = {adown_17->getOutput(0), repncspelan_13->getOutput(0)}; + auto cat_18 = network->addConcatenation(input_tensor_18, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 19 (P4/16-medium) + auto repncspelan_19 = RepNCSPELAN4(network, weightMap, *cat_18->getOutput(0), 768, 192, 192, 96, 3, "model.18"); + + // # avg-conv-down merge + // [-1, 1, ADown, [512]], + auto adown_20 = AConv(network, weightMap, *repncspelan_19->getOutput(0), 128, "model.19"); + // [[-1, 10], 1, Concat, [1]], # cat head P5 + ITensor* input_tensor_21[] = {adown_20->getOutput(0), sppelan_10->getOutput(0)}; + auto cat_21 = network->addConcatenation(input_tensor_21, 2); + + // # elan-2 block + auto repncspelan_22 = RepNCSPELAN4(network, weightMap, *cat_21->getOutput(0), 1024, 256, 256, 128, 1, "model.21"); + std::vector head; + if (!isConvert) { + // # elan-spp block + auto sppelan_23 = SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 512, 256, 128, "model.22"); + + // # up-concat merge + auto upsample_24 = network->addResize(*sppelan_23->getOutput(0)); + upsample_24->setResizeMode(ResizeMode::kNEAREST); + const float scales_24[] = {1.0, 2.0, 2.0}; + upsample_24->setScales(scales_24, 3); + // [[-1, 6], 1, Concat, [1]], # cat backbone P4 + ITensor* input_tensor_25[] = {upsample_24->getOutput(0), repncspelan_7->getOutput(0)}; + auto cat_25 = network->addConcatenation(input_tensor_25, 2); + + // # elan-2 block + auto repncspelan_26 = RepNCSPELAN4(network, weightMap, *cat_25->getOutput(0), 384, 192, 192, 96, 3, "model.25"); + + // # up-concat merge + auto upsample_27 = network->addResize(*repncspelan_26->getOutput(0)); + upsample_27->setResizeMode(ResizeMode::kNEAREST); + const float scales_27[] = {1.0, 2.0, 2.0}; + upsample_27->setScales(scales_27, 3); + // [[-1, 4], 1, Concat, [1]], # cat backbone P3 + ITensor* input_tensor_28[] = {upsample_27->getOutput(0), repncspelan_5->getOutput(0)}; + auto cat_28 = network->addConcatenation(input_tensor_28, 2); + + // # elan-2 block + auto repncspelan_29 = RepNCSPELAN4(network, weightMap, *cat_28->getOutput(0), 256, 128, 128, 64, 3, "model.28"); + head = DualDDetect(network, weightMap, std::vector{repncspelan_16, repncspelan_19, repncspelan_22}, + kNumClass, {128, 192, 256}, "model.29"); + } else { + head = DDetect(network, weightMap, std::vector{repncspelan_16, repncspelan_19, repncspelan_22}, + kNumClass, {128, 192, 256}, "model.22"); + } + + nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, head, false); + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + builder->setMaxBatchSize(kBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); + +#if defined(USE_FP16) + config->setFlag(nvinfer1::BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(nvinfer1::BuilderFlag::kINT8); + auto* calibrator = + new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + delete network; + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return serialized_model; +} +IHostMemory* build_engine_yolov9_m(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, + std::string& wts_name, bool isConvert) { + /* ------ Create the builder ------ */ + INetworkDefinition* network = builder->createNetworkV2(0U); + + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{3, kInputH, kInputW}); + assert(data); + std::map weightMap = loadWeights(wts_name); + + int begin = isConvert ? 0 : 1; + + // # conv down + // [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 + auto conv_1 = convBnSiLU(network, weightMap, *data, 32, 3, 2, 1, "model." + std::to_string(begin), 1); + begin += 1; + // # conv down + // [-1, 1, Conv, [128, 3, 2]], # 2-P2/4 + auto conv_2 = convBnSiLU(network, weightMap, *conv_1->getOutput(0), 64, 3, 2, 1, "model." + std::to_string(begin)); + begin += 1; + // # elan-1 block + // [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 3 + auto repncspelan_3 = RepNCSPELAN4(network, weightMap, *conv_2->getOutput(0), 128, 128, 128, 64, 1, + "model." + std::to_string(begin)); + begin += 1; + // # avg-conv down + // [-1, 1, ADown, [256]], # 4-P3/8 + auto adown_4 = AConv(network, weightMap, *repncspelan_3->getOutput(0), 240, "model." + std::to_string(begin)); + begin += 1; + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 5 + auto repncspelan_5 = RepNCSPELAN4(network, weightMap, *adown_4->getOutput(0), 256, 240, 240, 120, 1, + "model." + std::to_string(begin)); + begin += 1; + // # avg-conv down + // [-1, 1, ADown, [512]], # 6-P4/16 + auto adown_6 = AConv(network, weightMap, *repncspelan_5->getOutput(0), 360, "model." + std::to_string(begin)); + begin += 1; + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 7 + auto repncspelan_7 = RepNCSPELAN4(network, weightMap, *adown_6->getOutput(0), 512, 360, 360, 180, 1, + "model." + std::to_string(begin)); + begin += 1; + // # avg-conv down + // [-1, 1, ADown, [512]], # 8-P5/32 + auto adown_8 = AConv(network, weightMap, *repncspelan_7->getOutput(0), 480, "model." + std::to_string(begin)); + begin += 1; + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 9 + auto repncspelan_9 = RepNCSPELAN4(network, weightMap, *adown_8->getOutput(0), 512, 480, 480, 240, 1, + "model." + std::to_string(begin)); + begin += 1; + // # elan-spp block + // [-1, 1, SPPELAN, [512, 256]], # 10 + auto sppelan_10 = + SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 512, 480, 240, "model." + std::to_string(begin)); + begin += 3; + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_11 = network->addResize(*sppelan_10->getOutput(0)); + upsample_11->setResizeMode(ResizeMode::kNEAREST); + const float scales_11[] = {1.0, 2.0, 2.0}; + upsample_11->setScales(scales_11, 3); + // [[-1, 7], 1, Concat, [1]], # cat backbone P4 + ITensor* input_tensor_12[] = {upsample_11->getOutput(0), repncspelan_7->getOutput(0)}; + auto cat_12 = network->addConcatenation(input_tensor_12, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 13 + auto repncspelan_13 = RepNCSPELAN4(network, weightMap, *cat_12->getOutput(0), 1536, 360, 360, 180, 1, + "model." + std::to_string(begin)); + begin += 3; + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_14 = network->addResize(*repncspelan_13->getOutput(0)); + upsample_14->setResizeMode(ResizeMode::kNEAREST); + const float scales_14[] = {1.0, 2.0, 2.0}; + upsample_14->setScales(scales_14, 3); + // [[-1, 5], 1, Concat, [1]], # cat backbone P3 + ITensor* input_tensor_15[] = {upsample_14->getOutput(0), repncspelan_5->getOutput(0)}; + auto cat_15 = network->addConcatenation(input_tensor_15, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]], # 16 (P3/8-small) + auto repncspelan_16 = RepNCSPELAN4(network, weightMap, *cat_15->getOutput(0), 1024, 240, 240, 120, 1, + "model." + std::to_string(begin)); + begin += 1; + + // # avg-conv-down merge + // [-1, 1, ADown, [256]], + auto adown_17 = AConv(network, weightMap, *repncspelan_16->getOutput(0), 184, "model." + std::to_string(begin)); + begin += 2; + // [[-1, 13], 1, Concat, [1]], # cat head P4 + ITensor* input_tensor_18[] = {adown_17->getOutput(0), repncspelan_13->getOutput(0)}; + auto cat_18 = network->addConcatenation(input_tensor_18, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 19 (P4/16-medium) + auto repncspelan_19 = RepNCSPELAN4(network, weightMap, *cat_18->getOutput(0), 768, 360, 360, 180, 1, + "model." + std::to_string(begin)); + begin += 1; + + // # avg-conv-down merge + // [-1, 1, ADown, [512]], + auto adown_20 = AConv(network, weightMap, *repncspelan_19->getOutput(0), 240, "model." + std::to_string(begin)); + begin += 2; + // [[-1, 10], 1, Concat, [1]], # cat head P5 + ITensor* input_tensor_21[] = {adown_20->getOutput(0), sppelan_10->getOutput(0)}; + auto cat_21 = network->addConcatenation(input_tensor_21, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 22 (P5/32-large) + auto repncspelan_22 = RepNCSPELAN4(network, weightMap, *cat_21->getOutput(0), 1024, 480, 480, 240, 1, + "model." + std::to_string(begin)); + begin += 1; + std::vector head; + if (!isConvert) { + // # routing + // [5, 1, CBLinear, [[256]]], # 23 + auto cblinear_23 = CBLinear(network, weightMap, *repncspelan_5->getOutput(0), {240}, 1, 1, 0, 1, + "model." + std::to_string(begin)); + begin += 1; + // [7, 1, CBLinear, [[256, 512]]], # 24 + auto cblinear_24 = CBLinear(network, weightMap, *repncspelan_7->getOutput(0), {240, 360}, 1, 1, 0, 1, + "model." + std::to_string(begin)); + begin += 1; + // [9, 1, CBLinear, [[256, 512, 512]]], # 25 + auto cblinear_25 = CBLinear(network, weightMap, *repncspelan_9->getOutput(0), {240, 360, 480}, 1, 1, 0, 1, + "model." + std::to_string(begin)); + begin += 1; + + // # conv down + // [0, 1, Conv, [64, 3, 2]], # 26-P1/2 + auto conv_26 = convBnSiLU(network, weightMap, *data, 32, 3, 2, 1, "model." + std::to_string(begin), 1); + begin += 1; + + // # conv down + // [-1, 1, Conv, [128, 3, 2]], # 27-P2/4 + auto conv_27 = + convBnSiLU(network, weightMap, *conv_26->getOutput(0), 64, 3, 2, 1, "model." + std::to_string(begin)); + begin += 1; + + // # elan-1 block + // [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 28 + auto repncspelan_28 = RepNCSPELAN4(network, weightMap, *conv_27->getOutput(0), 128, 128, 128, 64, 1, + "model." + std::to_string(begin)); + begin += 1; + + // # avg-conv down fuse + // [-1, 1, ADown, [256]], # 29-P3/8 + auto adown_29 = AConv(network, weightMap, *repncspelan_28->getOutput(0), 240, "model." + std::to_string(begin)); + begin += 2; + // [[23, 24, 25, -1], 1, CBFuse, [[0, 0, 0]]], # 30 + auto cbfuse = CBFuse(network, {cblinear_23, cblinear_24, cblinear_25, std::vector{adown_29}}, + {0, 0, 0, 0}, {8, 16, 32, 8}); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 31 + auto repncspelan_31 = RepNCSPELAN4(network, weightMap, *cbfuse->getOutput(0), 256, 240, 240, 120, 1, + "model." + std::to_string(begin)); + begin += 1; + + // # avg-conv down fuse + // [-1, 1, ADown, [512]], # 32-P4/16 + auto adown_32 = AConv(network, weightMap, *repncspelan_31->getOutput(0), 360, "model." + std::to_string(begin)); + begin += 2; + // [[24, 25, -1], 1, CBFuse, [[1, 1]]], # 33 + auto cbfuse_33 = + CBFuse(network, {cblinear_24, cblinear_25, std::vector{adown_32}}, {1, 1, 0}, {16, 32, 16}); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 34 + auto repncspelan_34 = RepNCSPELAN4(network, weightMap, *cbfuse_33->getOutput(0), 512, 360, 360, 180, 1, + "model." + std::to_string(begin)); + begin += 1; + + // # avg-conv down fuse + // [-1, 1, ADown, [512]], # 35-P5/32 + auto adown_35 = AConv(network, weightMap, *repncspelan_34->getOutput(0), 480, "model." + std::to_string(begin)); + begin += 2; + + // [[25, -1], 1, CBFuse, [[2]]], # 36 + auto cbfuse_36 = CBFuse(network, {cblinear_25, std::vector{adown_35}}, {2, 0}, {32, 32}); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 37 + auto repncspelan_37 = RepNCSPELAN4(network, weightMap, *cbfuse_36->getOutput(0), 512, 480, 480, 240, 1, + "model." + std::to_string(begin)); + begin += 1; + + // # detection head + // # detect + // [[31, 34, 37, 16, 19, 22], 1, DualDDetect, [nc]], # DualDDetect(A3, A4, A5, P3, P4, P5) + head = DualDDetect(network, weightMap, std::vector{repncspelan_31, repncspelan_34, repncspelan_37}, + kNumClass, {240, 360, 480}, "model." + std::to_string(begin)); + } else { + // # detection head + // # detect + // [[16, 19, 22], 1, DDetect, [nc]], # DDetect(P3, P4, P5) + head = DDetect(network, weightMap, std::vector{repncspelan_16, repncspelan_19, repncspelan_22}, + kNumClass, {240, 360, 480}, "model." + std::to_string(begin)); + } + + nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, head, false); + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + builder->setMaxBatchSize(kBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); + +#if defined(USE_FP16) + config->setFlag(nvinfer1::BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(nvinfer1::BuilderFlag::kINT8); + auto* calibrator = + new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + delete network; + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return serialized_model; +} + +IHostMemory* build_engine_yolov9_c(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, + std::string& wts_name) { + /* ------ Create the builder ------ */ + INetworkDefinition* network = builder->createNetworkV2(0U); + + ITensor* data = network->addInput(kInputTensorName, dt, Dims3{3, kInputH, kInputW}); + assert(data); + std::map weightMap = loadWeights(wts_name); + + // # conv down + // [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 + auto conv_1 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.1", 1); + // # conv down + // [-1, 1, Conv, [128, 3, 2]], # 2-P2/4 + auto conv_2 = convBnSiLU(network, weightMap, *conv_1->getOutput(0), 128, 3, 2, 1, "model.2"); + // # elan-1 block + // [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 3 + auto repncspelan_3 = RepNCSPELAN4(network, weightMap, *conv_2->getOutput(0), 128, 256, 128, 64, 1, "model.3"); + // # avg-conv down + // [-1, 1, ADown, [256]], # 4-P3/8 + auto adown_4 = ADown(network, weightMap, *repncspelan_3->getOutput(0), 256, "model.4"); + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 5 + auto repncspelan_5 = RepNCSPELAN4(network, weightMap, *adown_4->getOutput(0), 256, 512, 256, 128, 1, "model.5"); + // # avg-conv down + // [-1, 1, ADown, [512]], # 6-P4/16 + auto adown_6 = ADown(network, weightMap, *repncspelan_5->getOutput(0), 512, "model.6"); + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 7 + auto repncspelan_7 = RepNCSPELAN4(network, weightMap, *adown_6->getOutput(0), 512, 512, 512, 256, 1, "model.7"); + // # avg-conv down + // [-1, 1, ADown, [512]], # 8-P5/32 + auto adown_8 = ADown(network, weightMap, *repncspelan_7->getOutput(0), 512, "model.8"); + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 9 + auto repncspelan_9 = RepNCSPELAN4(network, weightMap, *adown_8->getOutput(0), 512, 512, 512, 256, 1, "model.9"); + // # elan-spp block + // [-1, 1, SPPELAN, [512, 256]], # 10 + auto sppelan_10 = SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 512, 512, 256, "model.10"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_11 = network->addResize(*sppelan_10->getOutput(0)); + upsample_11->setResizeMode(ResizeMode::kNEAREST); + const float scales_11[] = {1.0, 2.0, 2.0}; + upsample_11->setScales(scales_11, 3); + // [[-1, 7], 1, Concat, [1]], # cat backbone P4 + ITensor* input_tensor_12[] = {upsample_11->getOutput(0), repncspelan_7->getOutput(0)}; + auto cat_12 = network->addConcatenation(input_tensor_12, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 13 + auto repncspelan_13 = RepNCSPELAN4(network, weightMap, *cat_12->getOutput(0), 1536, 512, 512, 256, 1, "model.13"); + + // # up-concat merge + // [-1, 1, nn.Upsample, [None, 2, 'nearest']], + auto upsample_14 = network->addResize(*repncspelan_13->getOutput(0)); + upsample_14->setResizeMode(ResizeMode::kNEAREST); + const float scales_14[] = {1.0, 2.0, 2.0}; + upsample_14->setScales(scales_14, 3); + // [[-1, 5], 1, Concat, [1]], # cat backbone P3 + ITensor* input_tensor_15[] = {upsample_14->getOutput(0), repncspelan_5->getOutput(0)}; + auto cat_15 = network->addConcatenation(input_tensor_15, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]], # 16 (P3/8-small) + auto repncspelan_16 = RepNCSPELAN4(network, weightMap, *cat_15->getOutput(0), 1024, 256, 256, 128, 1, "model.16"); + + // # avg-conv-down merge + // [-1, 1, ADown, [256]], + auto adown_17 = ADown(network, weightMap, *repncspelan_16->getOutput(0), 256, "model.17"); + // [[-1, 13], 1, Concat, [1]], # cat head P4 + ITensor* input_tensor_18[] = {adown_17->getOutput(0), repncspelan_13->getOutput(0)}; + auto cat_18 = network->addConcatenation(input_tensor_18, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 19 (P4/16-medium) + auto repncspelan_19 = RepNCSPELAN4(network, weightMap, *cat_18->getOutput(0), 768, 512, 512, 256, 1, "model.19"); + + // # avg-conv-down merge + // [-1, 1, ADown, [512]], + auto adown_20 = ADown(network, weightMap, *repncspelan_19->getOutput(0), 512, "model.20"); + // [[-1, 10], 1, Concat, [1]], # cat head P5 + ITensor* input_tensor_21[] = {adown_20->getOutput(0), sppelan_10->getOutput(0)}; + auto cat_21 = network->addConcatenation(input_tensor_21, 2); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 22 (P5/32-large) + auto repncspelan_22 = RepNCSPELAN4(network, weightMap, *cat_21->getOutput(0), 1024, 512, 512, 256, 1, "model.22"); + + // # multi-level reversible auxiliary branch + + // # routing + // [5, 1, CBLinear, [[256]]], # 23 + auto cblinear_23 = CBLinear(network, weightMap, *repncspelan_5->getOutput(0), {256}, 1, 1, 0, 1, "model.23"); + // [7, 1, CBLinear, [[256, 512]]], # 24 + auto cblinear_24 = CBLinear(network, weightMap, *repncspelan_7->getOutput(0), {256, 512}, 1, 1, 0, 1, "model.24"); + // [9, 1, CBLinear, [[256, 512, 512]]], # 25 + auto cblinear_25 = + CBLinear(network, weightMap, *repncspelan_9->getOutput(0), {256, 512, 512}, 1, 1, 0, 1, "model.25"); + + // # conv down + // [0, 1, Conv, [64, 3, 2]], # 26-P1/2 + auto conv_26 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.26", 1); + + // # conv down + // [-1, 1, Conv, [128, 3, 2]], # 27-P2/4 + auto conv_27 = convBnSiLU(network, weightMap, *conv_26->getOutput(0), 128, 3, 2, 1, "model.27"); + + // # elan-1 block + // [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 28 + auto repncspelan_28 = RepNCSPELAN4(network, weightMap, *conv_27->getOutput(0), 128, 256, 128, 64, 1, "model.28"); + + // # avg-conv down fuse + // [-1, 1, ADown, [256]], # 29-P3/8 + auto adown_29 = ADown(network, weightMap, *repncspelan_28->getOutput(0), 256, "model.29"); + // [[23, 24, 25, -1], 1, CBFuse, [[0, 0, 0]]], # 30 + auto cbfuse = CBFuse(network, {cblinear_23, cblinear_24, cblinear_25, std::vector{adown_29}}, {0, 0, 0, 0}, + {8, 16, 32, 8}); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 31 + auto repncspelan_31 = RepNCSPELAN4(network, weightMap, *cbfuse->getOutput(0), 256, 512, 256, 128, 1, "model.31"); + + // # avg-conv down fuse + // [-1, 1, ADown, [512]], # 32-P4/16 + auto adown_32 = ADown(network, weightMap, *repncspelan_31->getOutput(0), 512, "model.32"); + // [[24, 25, -1], 1, CBFuse, [[1, 1]]], # 33 + auto cbfuse_33 = + CBFuse(network, {cblinear_24, cblinear_25, std::vector{adown_32}}, {1, 1, 0}, {16, 32, 16}); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 34 + auto repncspelan_34 = RepNCSPELAN4(network, weightMap, *cbfuse_33->getOutput(0), 512, 512, 512, 256, 1, "model.34"); + + // # avg-conv down fuse + // [-1, 1, ADown, [512]], # 35-P5/32 + auto adown_35 = ADown(network, weightMap, *repncspelan_34->getOutput(0), 512, "model.35"); + + // [[25, -1], 1, CBFuse, [[2]]], # 36 + auto cbfuse_36 = CBFuse(network, {cblinear_25, std::vector{adown_35}}, {2, 0}, {32, 32}); + + // # elan-2 block + // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 37 + auto repncspelan_37 = RepNCSPELAN4(network, weightMap, *cbfuse_36->getOutput(0), 512, 512, 512, 256, 1, "model.37"); + + // # detection head + // # detect + // [[31, 34, 37, 16, 19, 22], 1, DualDDetect, [nc]], # DualDDetect(A3, A4, A5, P3, P4, P5) + auto dualddetect_38 = + DualDDetect(network, weightMap, std::vector{repncspelan_31, repncspelan_34, repncspelan_37}, + kNumClass, {512, 512, 512}, "model.38"); + + nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, dualddetect_38, false); + yolo->getOutput(0)->setName(kOutputTensorName); + network->markOutput(*yolo->getOutput(0)); + + builder->setMaxBatchSize(kBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); + +#if defined(USE_FP16) + config->setFlag(nvinfer1::BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(nvinfer1::BuilderFlag::kINT8); + auto* calibrator = + new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + + delete network; + + // Release host memory + for (auto& mem : weightMap) { + free((void*)(mem.second.values)); + } + + return serialized_model; +} + IHostMemory* build_engine_yolov9_e(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, std::string& wts_name) { /* ------ Create the builder ------ */ @@ -247,194 +964,9 @@ IHostMemory* build_engine_yolov9_e(unsigned int maxBatchSize, IBuilder* builder, return serialized_model; } -IHostMemory* build_engine_yolov9_c(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, - std::string& wts_name) { - /* ------ Create the builder ------ */ - INetworkDefinition* network = builder->createNetworkV2(0U); - ITensor* data = network->addInput(kInputTensorName, dt, Dims3{3, kInputH, kInputW}); - assert(data); - std::map weightMap = loadWeights(wts_name); - - // # conv down - // [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 - auto conv_1 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.1", 1); - // # conv down - // [-1, 1, Conv, [128, 3, 2]], # 2-P2/4 - auto conv_2 = convBnSiLU(network, weightMap, *conv_1->getOutput(0), 128, 3, 2, 1, "model.2"); - // # elan-1 block - // [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 3 - auto repncspelan_3 = RepNCSPELAN4(network, weightMap, *conv_2->getOutput(0), 128, 256, 128, 64, 1, "model.3"); - // # avg-conv down - // [-1, 1, ADown, [256]], # 4-P3/8 - auto adown_4 = ADown(network, weightMap, *repncspelan_3->getOutput(0), 256, "model.4"); - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 5 - auto repncspelan_5 = RepNCSPELAN4(network, weightMap, *adown_4->getOutput(0), 256, 512, 256, 128, 1, "model.5"); - // # avg-conv down - // [-1, 1, ADown, [512]], # 6-P4/16 - auto adown_6 = ADown(network, weightMap, *repncspelan_5->getOutput(0), 512, "model.6"); - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 7 - auto repncspelan_7 = RepNCSPELAN4(network, weightMap, *adown_6->getOutput(0), 512, 512, 512, 256, 1, "model.7"); - // # avg-conv down - // [-1, 1, ADown, [512]], # 8-P5/32 - auto adown_8 = ADown(network, weightMap, *repncspelan_7->getOutput(0), 512, "model.8"); - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 9 - auto repncspelan_9 = RepNCSPELAN4(network, weightMap, *adown_8->getOutput(0), 512, 512, 512, 256, 1, "model.9"); - // # elan-spp block - // [-1, 1, SPPELAN, [512, 256]], # 10 - auto sppelan_10 = SPPELAN(network, weightMap, *repncspelan_9->getOutput(0), 512, 512, 256, "model.10"); - - // # up-concat merge - // [-1, 1, nn.Upsample, [None, 2, 'nearest']], - auto upsample_11 = network->addResize(*sppelan_10->getOutput(0)); - upsample_11->setResizeMode(ResizeMode::kNEAREST); - const float scales_11[] = {1.0, 2.0, 2.0}; - upsample_11->setScales(scales_11, 3); - // [[-1, 7], 1, Concat, [1]], # cat backbone P4 - ITensor* input_tensor_12[] = {upsample_11->getOutput(0), repncspelan_7->getOutput(0)}; - auto cat_12 = network->addConcatenation(input_tensor_12, 2); - - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 13 - auto repncspelan_13 = RepNCSPELAN4(network, weightMap, *cat_12->getOutput(0), 1536, 512, 512, 256, 1, "model.13"); - - // # up-concat merge - // [-1, 1, nn.Upsample, [None, 2, 'nearest']], - auto upsample_14 = network->addResize(*repncspelan_13->getOutput(0)); - upsample_14->setResizeMode(ResizeMode::kNEAREST); - const float scales_14[] = {1.0, 2.0, 2.0}; - upsample_14->setScales(scales_14, 3); - // [[-1, 5], 1, Concat, [1]], # cat backbone P3 - ITensor* input_tensor_15[] = {upsample_14->getOutput(0), repncspelan_5->getOutput(0)}; - auto cat_15 = network->addConcatenation(input_tensor_15, 2); - - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]], # 16 (P3/8-small) - auto repncspelan_16 = RepNCSPELAN4(network, weightMap, *cat_15->getOutput(0), 1024, 256, 256, 128, 1, "model.16"); - - // # avg-conv-down merge - // [-1, 1, ADown, [256]], - auto adown_17 = ADown(network, weightMap, *repncspelan_16->getOutput(0), 256, "model.17"); - // [[-1, 13], 1, Concat, [1]], # cat head P4 - ITensor* input_tensor_18[] = {adown_17->getOutput(0), repncspelan_13->getOutput(0)}; - auto cat_18 = network->addConcatenation(input_tensor_18, 2); - - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 19 (P4/16-medium) - auto repncspelan_19 = RepNCSPELAN4(network, weightMap, *cat_18->getOutput(0), 768, 512, 512, 256, 1, "model.19"); - - // # avg-conv-down merge - // [-1, 1, ADown, [512]], - auto adown_20 = ADown(network, weightMap, *repncspelan_19->getOutput(0), 512, "model.20"); - // [[-1, 10], 1, Concat, [1]], # cat head P5 - ITensor* input_tensor_21[] = {adown_20->getOutput(0), sppelan_10->getOutput(0)}; - auto cat_21 = network->addConcatenation(input_tensor_21, 2); - - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 22 (P5/32-large) - auto repncspelan_22 = RepNCSPELAN4(network, weightMap, *cat_21->getOutput(0), 1024, 512, 512, 256, 1, "model.22"); - - // # multi-level reversible auxiliary branch - - // # routing - // [5, 1, CBLinear, [[256]]], # 23 - auto cblinear_23 = CBLinear(network, weightMap, *repncspelan_5->getOutput(0), {256}, 1, 1, 0, 1, "model.23"); - // [7, 1, CBLinear, [[256, 512]]], # 24 - auto cblinear_24 = CBLinear(network, weightMap, *repncspelan_7->getOutput(0), {256, 512}, 1, 1, 0, 1, "model.24"); - // [9, 1, CBLinear, [[256, 512, 512]]], # 25 - auto cblinear_25 = - CBLinear(network, weightMap, *repncspelan_9->getOutput(0), {256, 512, 512}, 1, 1, 0, 1, "model.25"); - - // # conv down - // [0, 1, Conv, [64, 3, 2]], # 26-P1/2 - auto conv_26 = convBnSiLU(network, weightMap, *data, 64, 3, 2, 1, "model.26", 1); - - // # conv down - // [-1, 1, Conv, [128, 3, 2]], # 27-P2/4 - auto conv_27 = convBnSiLU(network, weightMap, *conv_26->getOutput(0), 128, 3, 2, 1, "model.27"); - - // # elan-1 block - // [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]], # 28 - auto repncspelan_28 = RepNCSPELAN4(network, weightMap, *conv_27->getOutput(0), 128, 256, 128, 64, 1, "model.28"); - - // # avg-conv down fuse - // [-1, 1, ADown, [256]], # 29-P3/8 - auto adown_29 = ADown(network, weightMap, *repncspelan_28->getOutput(0), 256, "model.29"); - // [[23, 24, 25, -1], 1, CBFuse, [[0, 0, 0]]], # 30 - auto cbfuse = CBFuse(network, {cblinear_23, cblinear_24, cblinear_25, std::vector{adown_29}}, {0, 0, 0, 0}, - {8, 16, 32, 8}); - - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]], # 31 - auto repncspelan_31 = RepNCSPELAN4(network, weightMap, *cbfuse->getOutput(0), 256, 512, 256, 128, 1, "model.31"); - - // # avg-conv down fuse - // [-1, 1, ADown, [512]], # 32-P4/16 - auto adown_32 = ADown(network, weightMap, *repncspelan_31->getOutput(0), 512, "model.32"); - // [[24, 25, -1], 1, CBFuse, [[1, 1]]], # 33 - auto cbfuse_33 = - CBFuse(network, {cblinear_24, cblinear_25, std::vector{adown_32}}, {1, 1, 0}, {16, 32, 16}); - - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 34 - auto repncspelan_34 = RepNCSPELAN4(network, weightMap, *cbfuse_33->getOutput(0), 512, 512, 512, 256, 1, "model.34"); - - // # avg-conv down fuse - // [-1, 1, ADown, [512]], # 35-P5/32 - auto adown_35 = ADown(network, weightMap, *repncspelan_34->getOutput(0), 512, "model.35"); - - // [[25, -1], 1, CBFuse, [[2]]], # 36 - auto cbfuse_36 = CBFuse(network, {cblinear_25, std::vector{adown_35}}, {2, 0}, {32, 32}); - - // # elan-2 block - // [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]], # 37 - auto repncspelan_37 = RepNCSPELAN4(network, weightMap, *cbfuse_36->getOutput(0), 512, 512, 512, 256, 1, "model.37"); - - // # detection head - // # detect - // [[31, 34, 37, 16, 19, 22], 1, DualDDetect, [nc]], # DualDDetect(A3, A4, A5, P3, P4, P5) - auto dualddetect_38 = - DualDDetect(network, weightMap, std::vector{repncspelan_31, repncspelan_34, repncspelan_37}, - kNumClass, {512, 512, 512}, "model.38"); - - nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(network, dualddetect_38, false); - yolo->getOutput(0)->setName(kOutputTensorName); - network->markOutput(*yolo->getOutput(0)); - - builder->setMaxBatchSize(kBatchSize); - config->setMaxWorkspaceSize(16 * (1 << 20)); - -#if defined(USE_FP16) - config->setFlag(nvinfer1::BuilderFlag::kFP16); -#elif defined(USE_INT8) - std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; - assert(builder->platformHasFastInt8()); - config->setFlag(nvinfer1::BuilderFlag::kINT8); - auto* calibrator = - new Int8EntropyCalibrator2(1, kInputW, kInputH, gCalibTablePath, "int8calib.table", kInputTensorName); - config->setInt8Calibrator(calibrator); -#endif - - std::cout << "Building engine, please wait for a while..." << std::endl; - IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config); - std::cout << "Build engine successfully!" << std::endl; - - delete network; - - // Release host memory - for (auto& mem : weightMap) { - free((void*)(mem.second.values)); - } - - return serialized_model; -} - -nvinfer1::IHostMemory* build_engine_gelan_e(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, - nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, - std::string& wts_name) { +IHostMemory* build_engine_gelan_e(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, + std::string& wts_name) { /* ------ Create the builder ------ */ INetworkDefinition* network = builder->createNetworkV2(0U); @@ -625,9 +1157,8 @@ nvinfer1::IHostMemory* build_engine_gelan_e(unsigned int maxBatchSize, nvinfer1: return serialized_model; } -nvinfer1::IHostMemory* build_engine_gelan_c(unsigned int maxBatchSize, nvinfer1::IBuilder* builder, - nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, - std::string& wts_name) { +IHostMemory* build_engine_gelan_c(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, + std::string& wts_name) { /* ------ Create the builder ------ */ INetworkDefinition* network = builder->createNetworkV2(0U);