From 66d641814e683ca1cab5aca2a16140889961d758 Mon Sep 17 00:00:00 2001
From: JumpPandaer <834662983@qq.com>
Date: Wed, 7 Apr 2021 18:08:25 +0800
Subject: [PATCH] Update yolov5.cpp (#463)

* Update yolov5.cpp

implement upsample with IResizeLayer

* Update yolov5.cpp

repair mistake and test it
---
 yolov5/yolov5.cpp | 26 +++++++++++---------------
 1 file changed, 11 insertions(+), 15 deletions(-)

diff --git a/yolov5/yolov5.cpp b/yolov5/yolov5.cpp
index 9c95206..1e83ccc 100644
--- a/yolov5/yolov5.cpp
+++ b/yolov5/yolov5.cpp
@@ -62,26 +62,22 @@ ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilder
     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<float*>(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, 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, gw));
-    weightMap["deconv11"] = deconvwts11;
+    auto upsample11 = network->addResize(*conv10->getOutput(0));
+    assert(upsample11);
+    upsample11->setResizeMode(ResizeMode::kNEAREST);
+    upsample11->setOutputDimensions(bottleneck_csp6->getOutput(0)->getDimensions());
 
-    ITensor* inputTensors12[] = { deconv11->getOutput(0), bottleneck_csp6->getOutput(0) };
+    ITensor* inputTensors12[] = { upsample11->getOutput(0), bottleneck_csp6->getOutput(0) };
     auto cat12 = network->addConcatenation(inputTensors12, 2);
     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, 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, gw));
-    ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) };
+    auto upsample15 = network->addResize(*conv14->getOutput(0));
+    assert(upsample15);
+    upsample15->setResizeMode(ResizeMode::kNEAREST);
+    upsample15->setOutputDimensions(bottleneck_csp4->getOutput(0)->getDimensions());
+	
+    ITensor* inputTensors16[] = { upsample15->getOutput(0), bottleneck_csp4->getOutput(0) };
     auto cat16 = network->addConcatenation(inputTensors16, 2);
 
     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");