haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

yolov5s PANet update

Browse Source
This commit is contained in:
wang-xinyu 2020-06-23 21:19:10 +08:00
parent 271b152c30
commit bca85e29e8
1 changed files with 42 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
yolov5/yolov5s.cpp
View File

@ -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<float*>(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<float*>(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);