Add PP-YOLOE support

nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.cpp

@@ -5,40 +5,47 @@
 
 #include "shortcut_layer.h"
 
-nvinfer1::ILayer* shortcutLayer(
+nvinfer1::ITensor* shortcutLayer(
     int layerIdx,
+    std::string mode,
     std::string activation,
     std::string inputVol,
     std::string shortcutVol,
+    std::map<std::string, std::string>& block,
     nvinfer1::ITensor* input,
-    nvinfer1::ITensor* shortcutTensor,
+    nvinfer1::ITensor* shortcutInput,
     nvinfer1::INetworkDefinition* network)
 {
-    nvinfer1::ILayer* output;
-    nvinfer1::ITensor* outputTensor;
+    nvinfer1::ITensor* output;
 
-    if (inputVol != shortcutVol)
+    assert(block.at("type") == "shortcut");
+
+    nvinfer1::ElementWiseOperation operation = nvinfer1::ElementWiseOperation::kSUM;
+
+    if (mode == "mul")
+        operation = nvinfer1::ElementWiseOperation::kPROD;
+
+    if (mode == "add" && inputVol != shortcutVol)
     {
-        nvinfer1::ISliceLayer* sl = network->addSlice(
-            *shortcutTensor,
-            nvinfer1::Dims3{0, 0, 0},
-            input->getDimensions(),
-            nvinfer1::Dims3{1, 1, 1});
-        assert(sl != nullptr);
-        outputTensor = sl->getOutput(0);
-        assert(outputTensor != nullptr);
-    } else 
+        nvinfer1::ISliceLayer* slice = network->addSlice(
+            *shortcutInput, nvinfer1::Dims{3, {0, 0, 0}}, input->getDimensions(), nvinfer1::Dims{3, {1, 1, 1}});
+        assert(slice != nullptr);
+        std::string sliceLayerName = "slice_" + std::to_string(layerIdx);
+        slice->setName(sliceLayerName.c_str());
+        output = slice->getOutput(0);
+    }
+    else 
     {
-        outputTensor = shortcutTensor;
-        assert(outputTensor != nullptr);
+        output = shortcutInput;
     }
 
-    nvinfer1::IElementWiseLayer* ew = network->addElementWise(
-        *input, *outputTensor,
-        nvinfer1::ElementWiseOperation::kSUM);
-    assert(ew != nullptr);
+    nvinfer1::IElementWiseLayer* shortcut = network->addElementWise(*input, *output, operation);
+    assert(shortcut != nullptr);
+    std::string shortcutLayerName = "shortcut_" + std::to_string(layerIdx);
+    shortcut->setName(shortcutLayerName.c_str());
+    output = shortcut->getOutput(0);
 
-    output = activationLayer(layerIdx, activation, ew, ew->getOutput(0), network);
+    output = activationLayer(layerIdx, activation, output, network);
     assert(output != nullptr);
 
     return output;