From bca85e29e89cfadd8fc233fb37eea10447999be0 Mon Sep 17 00:00:00 2001 From: wang-xinyu Date: Tue, 23 Jun 2020 21:19:10 +0800 Subject: [PATCH] yolov5s PANet update --- yolov5/yolov5s.cpp | 70 +++++++++++++++++++++++++++------------------- 1 file changed, 42 insertions(+), 28 deletions(-) diff --git a/yolov5/yolov5s.cpp b/yolov5/yolov5s.cpp index cc2f951..9fc0932 100644 --- a/yolov5/yolov5s.cpp +++ b/yolov5/yolov5s.cpp @@ -33,48 +33,62 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilder // yolov5 backbone auto focus0 = focus(network, weightMap, *data, 3, 32, 3, "model.0"); auto conv1 = convBnLeaky(network, weightMap, *focus0->getOutput(0), 64, 3, 2, 1, "model.1"); - auto bottleneck2 = bottleneck(network, weightMap, *conv1->getOutput(0), 64, 64, true, 1, 0.5, "model.2"); - auto conv3 = convBnLeaky(network, weightMap, *bottleneck2->getOutput(0), 128, 3, 2, 1, "model.3"); + auto bottleneck_CSP2 = bottleneckCSP(network, weightMap, *conv1->getOutput(0), 64, 64, 1, true, 1, 0.5, "model.2"); + auto conv3 = convBnLeaky(network, weightMap, *bottleneck_CSP2->getOutput(0), 128, 3, 2, 1, "model.3"); auto bottleneck_csp4 = bottleneckCSP(network, weightMap, *conv3->getOutput(0), 128, 128, 3, true, 1, 0.5, "model.4"); auto conv5 = convBnLeaky(network, weightMap, *bottleneck_csp4->getOutput(0), 256, 3, 2, 1, "model.5"); auto bottleneck_csp6 = bottleneckCSP(network, weightMap, *conv5->getOutput(0), 256, 256, 3, true, 1, 0.5, "model.6"); auto conv7 = convBnLeaky(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"); - auto bottleneck_csp9 = bottleneckCSP(network, weightMap, *spp8->getOutput(0), 512, 512, 2, true, 1, 0.5, "model.9"); - // yolov5 head - auto bottleneck_csp10 = bottleneckCSP(network, weightMap, *bottleneck_csp9->getOutput(0), 512, 512, 1, false, 1, 0.5, "model.10"); - IConvolutionLayer* conv11 = network->addConvolutionNd(*bottleneck_csp10->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["model.11.weight"], weightMap["model.11.bias"]); - float *deval = reinterpret_cast(malloc(sizeof(float) * 512 * 2 * 2)); - for (int i = 0; i < 512 * 2 * 2; i++) { + // yolov5 head + auto bottleneck_csp9 = bottleneckCSP(network, weightMap, *spp8->getOutput(0), 512, 512, 1, false, 1, 0.5, "model.9"); + auto conv10 = convBnLeaky(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 deconvwts12{DataType::kFLOAT, deval, 512 * 2 * 2}; - IDeconvolutionLayer* deconv12 = network->addDeconvolutionNd(*bottleneck_csp10->getOutput(0), 512, DimsHW{2, 2}, deconvwts12, emptywts); - deconv12->setStrideNd(DimsHW{2, 2}); - deconv12->setNbGroups(512); - weightMap["deconv12"] = deconvwts12; + 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* inputTensors13[] = {deconv12->getOutput(0), bottleneck_csp6->getOutput(0)}; - auto cat13 = network->addConcatenation(inputTensors13, 2); - auto conv14 = convBnLeaky(network, weightMap, *cat13->getOutput(0), 256, 1, 1, 1, "model.14"); - auto bottleneck_csp15 = bottleneckCSP(network, weightMap, *conv14->getOutput(0), 256, 256, 1, false, 1, 0.5, "model.15"); - IConvolutionLayer* conv16 = network->addConvolutionNd(*bottleneck_csp15->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["model.16.weight"], weightMap["model.16.bias"]); + ITensor* inputTensors12[] = {deconv11->getOutput(0), bottleneck_csp6->getOutput(0)}; + auto cat12 = network->addConcatenation(inputTensors12, 2); + auto bottleneck_csp13 = bottleneckCSP(network, weightMap, *cat12->getOutput(0), 512, 256, 1, false, 1, 0.5, "model.13"); + auto conv14 = convBnLeaky(network, weightMap, *bottleneck_csp13->getOutput(0), 128, 1, 1, 1, "model.14"); + std::cout << "conv14 ----" << std::endl; - Weights deconvwts17{DataType::kFLOAT, deval, 256 * 2 * 2}; - IDeconvolutionLayer* deconv17 = network->addDeconvolutionNd(*bottleneck_csp15->getOutput(0), 256, DimsHW{2, 2}, deconvwts17, emptywts); - deconv17->setStrideNd(DimsHW{2, 2}); - deconv17->setNbGroups(256); - ITensor* inputTensors18[] = {deconv17->getOutput(0), bottleneck_csp4->getOutput(0)}; - auto cat18 = network->addConcatenation(inputTensors18, 2); - auto conv19 = convBnLeaky(network, weightMap, *cat18->getOutput(0), 128, 1, 1, 1, "model.19"); - auto bottleneck_csp20 = bottleneckCSP(network, weightMap, *conv19->getOutput(0), 128, 128, 1, false, 1, 0.5, "model.20"); - IConvolutionLayer* conv21 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["model.21.weight"], weightMap["model.21.bias"]); + 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 = bottleneckCSP(network, weightMap, *cat16->getOutput(0), 256, 128, 1, false, 1, 0.5, "model.17"); + std::cout << "conv18 ----" << std::endl; + IConvolutionLayer* conv18 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["model.18.weight"], weightMap["model.18.bias"]); + + auto conv19 = convBnLeaky(network, weightMap, *bottleneck_csp17->getOutput(0), 128, 3, 2, 1, "model.19"); + ITensor* inputTensors20[] = {conv19->getOutput(0), conv14->getOutput(0)}; + auto cat20 = network->addConcatenation(inputTensors20, 2); + auto bottleneck_csp21 = bottleneckCSP(network, weightMap, *cat20->getOutput(0), 256, 256, 1, false, 1, 0.5, "model.21"); + std::cout << "conv22 ----" << std::endl; + IConvolutionLayer* conv22 = network->addConvolutionNd(*bottleneck_csp21->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["model.22.weight"], weightMap["model.22.bias"]); + + auto conv23 = convBnLeaky(network, weightMap, *bottleneck_csp21->getOutput(0), 256, 3, 2, 1, "model.23"); + ITensor* inputTensors24[] = {conv23->getOutput(0), conv10->getOutput(0)}; + auto cat24 = network->addConcatenation(inputTensors24, 2); + auto bottleneck_csp25 = bottleneckCSP(network, weightMap, *cat24->getOutput(0), 512, 512, 1, false, 1, 0.5, "model.25"); + IConvolutionLayer* conv26 = network->addConvolutionNd(*bottleneck_csp25->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{1, 1}, weightMap["model.26.weight"], weightMap["model.26.bias"]); auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); const PluginFieldCollection* pluginData = creator->getFieldNames(); IPluginV2 *pluginObj = creator->createPlugin("yololayer", pluginData); - ITensor* inputTensors_yolo[] = {conv11->getOutput(0), conv16->getOutput(0), conv21->getOutput(0)}; + ITensor* inputTensors_yolo[] = {conv26->getOutput(0), conv22->getOutput(0), conv18->getOutput(0)}; auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj); yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME);