DeepStream 7.1 + Fixes + New model output format

nvdsinfer_custom_impl_Yolo/layers/activation_layer.cpp

@@ -14,8 +14,9 @@ activationLayer(int layerIdx, std::string activation, nvinfer1::ITensor* input,
 {
   nvinfer1::ITensor* output;
 
-  if (activation == "linear")
+  if (activation == "linear") {
     output = input;
+  }
   else if (activation == "relu") {
     nvinfer1::IActivationLayer* relu = network->addActivation(*input, nvinfer1::ActivationType::kRELU);
     assert(relu != nullptr);

nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.cpp

@@ -21,6 +21,11 @@ batchnormLayer(int layerIdx, std::map<std::string, std::string>& block, std::vec
   int filters = std::stoi(block.at("filters"));
   std::string activation = block.at("activation");
 
+  float eps = 1.0e-5;
+  if (block.find("eps") != block.end()) {
+    eps = std::stof(block.at("eps"));
+  }
+
   std::vector<float> bnBiases;
   std::vector<float> bnWeights;
   std::vector<float> bnRunningMean;
@@ -39,7 +44,7 @@ batchnormLayer(int layerIdx, std::map<std::string, std::string>& block, std::vec
     ++weightPtr;
   }
   for (int i = 0; i < filters; ++i) {
-    bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5));
+    bnRunningVar.push_back(sqrt(weights[weightPtr] + eps));
     ++weightPtr;
   }
 
@@ -47,18 +52,25 @@ batchnormLayer(int layerIdx, std::map<std::string, std::string>& block, std::vec
   nvinfer1::Weights shift {nvinfer1::DataType::kFLOAT, nullptr, size};
   nvinfer1::Weights scale {nvinfer1::DataType::kFLOAT, nullptr, size};
   nvinfer1::Weights power {nvinfer1::DataType::kFLOAT, nullptr, size};
+
   float* shiftWt = new float[size];
-  for (int i = 0; i < size; ++i)
+  for (int i = 0; i < size; ++i) {
     shiftWt[i] = bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i));
+  }
   shift.values = shiftWt;
+
   float* scaleWt = new float[size];
-  for (int i = 0; i < size; ++i)
+  for (int i = 0; i < size; ++i) {
     scaleWt[i] = bnWeights.at(i) / bnRunningVar[i];
+  }
   scale.values = scaleWt;
+
   float* powerWt = new float[size];
-  for (int i = 0; i < size; ++i)
+  for (int i = 0; i < size; ++i) {
     powerWt[i] = 1.0;
+  }
   power.values = powerWt;
+
   trtWeights.push_back(shift);
   trtWeights.push_back(scale);
   trtWeights.push_back(power);

nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.cpp

@@ -15,7 +15,7 @@ convolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, std:
 {
   nvinfer1::ITensor* output;
 
-  assert(block.at("type") == "convolutional" || block.at("type") == "c2f");
+  assert(block.at("type") == "conv" || block.at("type") == "convolutional");
   assert(block.find("filters") != block.end());
   assert(block.find("pad") != block.end());
   assert(block.find("size") != block.end());
@@ -28,27 +28,35 @@ convolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, std:
   std::string activation = block.at("activation");
   int bias = filters;
 
-  bool batchNormalize = false;
+  int batchNormalize = 0;
+  float eps = 1.0e-5;
   if (block.find("batch_normalize") != block.end()) {
     bias = 0;
     batchNormalize = (block.at("batch_normalize") == "1");
+    if (block.find("eps") != block.end()) {
+      eps = std::stof(block.at("eps"));
+    }
   }
 
   if (block.find("bias") != block.end()) {
     bias = std::stoi(block.at("bias"));
-    if (bias == 1)
+    if (bias == 1) {
       bias = filters;
+    }
   }
 
   int groups = 1;
-  if (block.find("groups") != block.end())
+  if (block.find("groups") != block.end()) {
     groups = std::stoi(block.at("groups"));
+  }
 
   int pad;
-  if (padding)
+  if (padding) {
     pad = (kernelSize - 1) / 2;
-  else
+  }
+  else {
     pad = 0;
+  }
 
   int size = filters * inputChannels * kernelSize * kernelSize / groups;
   std::vector<float> bnBiases;
@@ -58,7 +66,7 @@ convolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, std:
   nvinfer1::Weights convWt {nvinfer1::DataType::kFLOAT, nullptr, size};
   nvinfer1::Weights convBias {nvinfer1::DataType::kFLOAT, nullptr, bias};
 
-  if (batchNormalize == false) {
+  if (batchNormalize == 0) {
     float* val;
     if (bias != 0) {
       val = new float[filters];
@@ -91,7 +99,7 @@ convolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, std:
       ++weightPtr;
     }
     for (int i = 0; i < filters; ++i) {
-      bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5));
+      bnRunningVar.push_back(sqrt(weights[weightPtr] + eps));
       ++weightPtr;
     }
     float* val;
@@ -110,40 +118,49 @@ convolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, std:
     }
     convWt.values = val;
     trtWeights.push_back(convWt);
-    if (bias != 0)
+    if (bias != 0) {
       trtWeights.push_back(convBias);
+    }
   }
 
-  nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(*input, filters, nvinfer1::Dims{2, {kernelSize, kernelSize}},
-      convWt, convBias);
+  nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd(*input, filters,
+      nvinfer1::Dims{2, {kernelSize, kernelSize}}, convWt, convBias);
   assert(conv != nullptr);
   std::string convLayerName = "conv_" + layerName + std::to_string(layerIdx);
   conv->setName(convLayerName.c_str());
   conv->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
   conv->setPaddingNd(nvinfer1::Dims{2, {pad, pad}});
 
-  if (block.find("groups") != block.end())
+  if (block.find("groups") != block.end()) {
     conv->setNbGroups(groups);
+  }
 
   output = conv->getOutput(0);
 
-  if (batchNormalize == true) {
+  if (batchNormalize == 1) {
     size = filters;
     nvinfer1::Weights shift {nvinfer1::DataType::kFLOAT, nullptr, size};
     nvinfer1::Weights scale {nvinfer1::DataType::kFLOAT, nullptr, size};
     nvinfer1::Weights power {nvinfer1::DataType::kFLOAT, nullptr, size};
+
     float* shiftWt = new float[size];
-    for (int i = 0; i < size; ++i)
+    for (int i = 0; i < size; ++i) {
       shiftWt[i] = bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i));
+    }
     shift.values = shiftWt;
+
     float* scaleWt = new float[size];
-    for (int i = 0; i < size; ++i)
+    for (int i = 0; i < size; ++i) {
       scaleWt[i] = bnWeights.at(i) / bnRunningVar[i];
+    }
     scale.values = scaleWt;
+
     float* powerWt = new float[size];
-    for (int i = 0; i < size; ++i)
+    for (int i = 0; i < size; ++i) {
       powerWt[i] = 1.0;
+    }
     power.values = powerWt;
+
     trtWeights.push_back(shift);
     trtWeights.push_back(scale);
     trtWeights.push_back(power);

nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.h

@@ -13,8 +13,8 @@
 
 #include "activation_layer.h"
 
-nvinfer1::ITensor* convolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, std::vector<float>& weights,
-    std::vector<nvinfer1::Weights>& trtWeights, int& weightPtr, int& inputChannels, nvinfer1::ITensor* input,
-    nvinfer1::INetworkDefinition* network, std::string layerName = "");
+nvinfer1::ITensor* convolutionalLayer(int layerIdx, std::map<std::string, std::string>& block,
+    std::vector<float>& weights, std::vector<nvinfer1::Weights>& trtWeights, int& weightPtr, int& inputChannels,
+    nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network, std::string layerName = "");
 
 #endif

nvdsinfer_custom_impl_Yolo/layers/deconvolutional_layer.cpp

@@ -6,6 +6,7 @@
 #include "deconvolutional_layer.h"
 
 #include <cassert>
+#include <math.h>
 
 nvinfer1::ITensor*
 deconvolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, std::vector<float>& weights,
@@ -14,7 +15,7 @@ deconvolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, st
 {
   nvinfer1::ITensor* output;
 
-  assert(block.at("type") == "deconvolutional");
+  assert(block.at("type") == "deconv" || block.at("type") == "deconvolutional");
   assert(block.find("filters") != block.end());
   assert(block.find("pad") != block.end());
   assert(block.find("size") != block.end());
@@ -24,20 +25,38 @@ deconvolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, st
   int padding = std::stoi(block.at("pad"));
   int kernelSize = std::stoi(block.at("size"));
   int stride = std::stoi(block.at("stride"));
+  std::string activation = block.at("activation");
   int bias = filters;
 
-  int groups = 1;
-  if (block.find("groups") != block.end())
-    groups = std::stoi(block.at("groups"));
+  int batchNormalize = 0;
+  float eps = 1.0e-5;
+  if (block.find("batch_normalize") != block.end()) {
+    bias = 0;
+    batchNormalize = (block.at("batch_normalize") == "1");
+    if (block.find("eps") != block.end()) {
+      eps = std::stof(block.at("eps"));
+    }
+  }
 
-  if (block.find("bias") != block.end())
+  if (block.find("bias") != block.end()) {
     bias = std::stoi(block.at("bias"));
+    if (bias == 1) {
+      bias = filters;
+    }
+  }
+
+  int groups = 1;
+  if (block.find("groups") != block.end()) {
+    groups = std::stoi(block.at("groups"));
+  }
 
   int pad;
-  if (padding)
+  if (padding) {
     pad = (kernelSize - 1) / 2;
-  else
+  }
+  else {
     pad = 0;
+  }
 
   int size = filters * inputChannels * kernelSize * kernelSize / groups;
   std::vector<float> bnBiases;
@@ -47,23 +66,62 @@ deconvolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, st
   nvinfer1::Weights convWt {nvinfer1::DataType::kFLOAT, nullptr, size};
   nvinfer1::Weights convBias {nvinfer1::DataType::kFLOAT, nullptr, bias};
 
-  float* val;
-  if (bias != 0) {
-    val = new float[filters];
-    for (int i = 0; i < filters; ++i) {
+  if (batchNormalize == 0) {
+    float* val;
+    if (bias != 0) {
+      val = new float[filters];
+      for (int i = 0; i < filters; ++i) {
+          val[i] = weights[weightPtr];
+          ++weightPtr;
+      }
+      convBias.values = val;
+      trtWeights.push_back(convBias);
+    }
+    val = new float[size];
+    for (int i = 0; i < size; ++i) {
         val[i] = weights[weightPtr];
         ++weightPtr;
     }
-    convBias.values = val;
-    trtWeights.push_back(convBias);
+    convWt.values = val;
+    trtWeights.push_back(convWt);
   }
-  val = new float[size];
-  for (int i = 0; i < size; ++i) {
+  else {
+    for (int i = 0; i < filters; ++i) {
+      bnBiases.push_back(weights[weightPtr]);
+      ++weightPtr;
+    }
+    for (int i = 0; i < filters; ++i) {
+      bnWeights.push_back(weights[weightPtr]);
+      ++weightPtr;
+    }
+    for (int i = 0; i < filters; ++i) {
+      bnRunningMean.push_back(weights[weightPtr]);
+      ++weightPtr;
+    }
+    for (int i = 0; i < filters; ++i) {
+      bnRunningVar.push_back(sqrt(weights[weightPtr] + eps));
+      ++weightPtr;
+    }
+    float* val;
+    if (bias != 0) {
+      val = new float[filters];
+      for (int i = 0; i < filters; ++i) {
+        val[i] = weights[weightPtr];
+        ++weightPtr;
+      }
+      convBias.values = val;
+    }
+    val = new float[size];
+    for (int i = 0; i < size; ++i) {
       val[i] = weights[weightPtr];
       ++weightPtr;
+    }
+    convWt.values = val;
+    trtWeights.push_back(convWt);
+    if (bias != 0) {
+      trtWeights.push_back(convBias);
+    }
   }
-  convWt.values = val;
-  trtWeights.push_back(convWt);
 
   nvinfer1::IDeconvolutionLayer* conv = network->addDeconvolutionNd(*input, filters,
       nvinfer1::Dims{2, {kernelSize, kernelSize}}, convWt, convBias);
@@ -73,10 +131,49 @@ deconvolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, st
   conv->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
   conv->setPaddingNd(nvinfer1::Dims{2, {pad, pad}});
 
-  if (block.find("groups") != block.end())
+  if (block.find("groups") != block.end()) {
     conv->setNbGroups(groups);
+  }
 
   output = conv->getOutput(0);
 
+  if (batchNormalize == 1) {
+    size = filters;
+    nvinfer1::Weights shift {nvinfer1::DataType::kFLOAT, nullptr, size};
+    nvinfer1::Weights scale {nvinfer1::DataType::kFLOAT, nullptr, size};
+    nvinfer1::Weights power {nvinfer1::DataType::kFLOAT, nullptr, size};
+
+    float* shiftWt = new float[size];
+    for (int i = 0; i < size; ++i) {
+      shiftWt[i] = bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i));
+    }
+    shift.values = shiftWt;
+
+    float* scaleWt = new float[size];
+    for (int i = 0; i < size; ++i) {
+      scaleWt[i] = bnWeights.at(i) / bnRunningVar[i];
+    }
+    scale.values = scaleWt;
+
+    float* powerWt = new float[size];
+    for (int i = 0; i < size; ++i) {
+      powerWt[i] = 1.0;
+    }
+    power.values = powerWt;
+
+    trtWeights.push_back(shift);
+    trtWeights.push_back(scale);
+    trtWeights.push_back(power);
+
+    nvinfer1::IScaleLayer* batchnorm = network->addScale(*output, nvinfer1::ScaleMode::kCHANNEL, shift, scale, power);
+    assert(batchnorm != nullptr);
+    std::string batchnormLayerName = "batchnorm_" + layerName + std::to_string(layerIdx);
+    batchnorm->setName(batchnormLayerName.c_str());
+    output = batchnorm->getOutput(0);
+  }
+
+  output = activationLayer(layerIdx, activation, output, network, layerName);
+  assert(output != nullptr);
+
   return output;
 }

nvdsinfer_custom_impl_Yolo/layers/deconvolutional_layer.h

@@ -8,12 +8,13 @@
 
 #include <map>
 #include <vector>
-#include <string>
 
 #include "NvInfer.h"
 
-nvinfer1::ITensor* deconvolutionalLayer(int layerIdx, std::map<std::string, std::string>& block, std::vector<float>& weights,
-    std::vector<nvinfer1::Weights>& trtWeights, int& weightPtr, int& inputChannels, nvinfer1::ITensor* input,
-    nvinfer1::INetworkDefinition* network, std::string layerName = "");
+#include "activation_layer.h"
+
+nvinfer1::ITensor* deconvolutionalLayer(int layerIdx, std::map<std::string, std::string>& block,
+    std::vector<float>& weights, std::vector<nvinfer1::Weights>& trtWeights, int& weightPtr, int& inputChannels,
+    nvinfer1::ITensor* input, nvinfer1::INetworkDefinition* network, std::string layerName = "");
 
 #endif

nvdsinfer_custom_impl_Yolo/layers/reorg_layer.cpp

@@ -10,7 +10,7 @@
 
 nvinfer1::ITensor*
 reorgLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1::ITensor* input,
-    nvinfer1::INetworkDefinition* network, uint batchSize)
+    nvinfer1::INetworkDefinition* network)
 {
   nvinfer1::ITensor* output;
 
@@ -35,17 +35,17 @@ reorgLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1::IT
     nvinfer1::Dims sizeAll = {4, {inputDims.d[0], inputDims.d[1], inputDims.d[2] / stride, inputDims.d[3] / stride}};
     nvinfer1::Dims strideAll = {4, {1, 1, stride, stride}};
 
-    nvinfer1::ITensor* slice1 = sliceLayer(layerIdx, name1, input, start1, sizeAll, strideAll, network, batchSize);
-    assert(output != nullptr);
+    nvinfer1::ITensor* slice1 = sliceLayer(layerIdx, name1, input, start1, sizeAll, strideAll, network);
+    assert(slice1 != nullptr);
 
-    nvinfer1::ITensor* slice2 = sliceLayer(layerIdx, name2, input, start2, sizeAll, strideAll, network, batchSize);
-    assert(output != nullptr);
+    nvinfer1::ITensor* slice2 = sliceLayer(layerIdx, name2, input, start2, sizeAll, strideAll, network);
+    assert(slice2 != nullptr);
 
-    nvinfer1::ITensor* slice3 = sliceLayer(layerIdx, name3, input, start3, sizeAll, strideAll, network, batchSize);
-    assert(output != nullptr);
+    nvinfer1::ITensor* slice3 = sliceLayer(layerIdx, name3, input, start3, sizeAll, strideAll, network);
+    assert(slice3 != nullptr);
 
-    nvinfer1::ITensor* slice4 = sliceLayer(layerIdx, name4, input, start4, sizeAll, strideAll, network, batchSize);
-    assert(output != nullptr);
+    nvinfer1::ITensor* slice4 = sliceLayer(layerIdx, name4, input, start4, sizeAll, strideAll, network);
+    assert(slice4 != nullptr);
 
     std::vector<nvinfer1::ITensor*> concatInputs;
     concatInputs.push_back(slice1);

nvdsinfer_custom_impl_Yolo/layers/reorg_layer.h

@@ -14,6 +14,6 @@
 #include "slice_layer.h"
 
 nvinfer1::ITensor* reorgLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1::ITensor* input,
-    nvinfer1::INetworkDefinition* network, uint batchSize);
+    nvinfer1::INetworkDefinition* network);
 
 #endif

nvdsinfer_custom_impl_Yolo/layers/route_layer.cpp

@@ -7,7 +7,7 @@
 
 nvinfer1::ITensor*
 routeLayer(int layerIdx, std::string& layers, std::map<std::string, std::string>& block,
-    std::vector<nvinfer1::ITensor*> tensorOutputs, nvinfer1::INetworkDefinition* network, uint batchSize)
+    std::vector<nvinfer1::ITensor*> tensorOutputs, nvinfer1::INetworkDefinition* network)
 {
   nvinfer1::ITensor* output;
 
@@ -49,7 +49,6 @@ routeLayer(int layerIdx, std::string& layers, std::map<std::string, std::string>
     int axis = 1;
     if (block.find("axis") != block.end()) {
       axis += std::stoi(block.at("axis"));
-      std::cout << axis << std::endl;
     }
     if (axis < 0) {
       axis += concatInputs[0]->getDimensions().nbDims;
@@ -75,7 +74,7 @@ routeLayer(int layerIdx, std::string& layers, std::map<std::string, std::string>
     nvinfer1::Dims size = {4, {prevTensorDims.d[0], channelSlice, prevTensorDims.d[2], prevTensorDims.d[3]}};
     nvinfer1::Dims stride = {4, {1, 1, 1, 1}};
 
-    output = sliceLayer(layerIdx, name, output, start, size, stride, network, batchSize);
+    output = sliceLayer(layerIdx, name, output, start, size, stride, network);
     assert(output != nullptr);
   }
 

nvdsinfer_custom_impl_Yolo/layers/route_layer.h

@@ -11,6 +11,6 @@
 #include "slice_layer.h"
 
 nvinfer1::ITensor* routeLayer(int layerIdx, std::string& layers, std::map<std::string, std::string>& block,
-    std::vector<nvinfer1::ITensor*> tensorOutputs, nvinfer1::INetworkDefinition* network, uint batchSize);
+    std::vector<nvinfer1::ITensor*> tensorOutputs, nvinfer1::INetworkDefinition* network);
 
 #endif

nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.cpp

@@ -10,7 +10,7 @@
 nvinfer1::ITensor*
 shortcutLayer(int layerIdx, std::string activation, std::string inputVol, std::string shortcutVol,
     std::map<std::string, std::string>& block, nvinfer1::ITensor* input, nvinfer1::ITensor* shortcutInput,
-    nvinfer1::INetworkDefinition* network, uint batchSize)
+    nvinfer1::INetworkDefinition* network)
 {
   nvinfer1::ITensor* output;
 
@@ -20,15 +20,17 @@ shortcutLayer(int layerIdx, std::string activation, std::string inputVol, std::s
     std::string name = "slice";
     nvinfer1::Dims start = {4, {0, 0, 0, 0}};
     nvinfer1::Dims size = input->getDimensions();
-    nvinfer1::Dims stride = nvinfer1::Dims{4, {1, 1, 1, 1}};
+    nvinfer1::Dims stride = {4, {1, 1, 1, 1}};
 
-    output = sliceLayer(layerIdx, name, shortcutInput, start, size, stride, network, batchSize);
+    output = sliceLayer(layerIdx, name, shortcutInput, start, size, stride, network);
     assert(output != nullptr);
   }
-  else
+  else {
     output = shortcutInput;
+  }
 
-  nvinfer1::IElementWiseLayer* shortcut = network->addElementWise(*input, *output, nvinfer1::ElementWiseOperation::kSUM);
+  nvinfer1::IElementWiseLayer* shortcut = network->addElementWise(*input, *output,
+      nvinfer1::ElementWiseOperation::kSUM);
   assert(shortcut != nullptr);
   std::string shortcutLayerName = "shortcut_" + std::to_string(layerIdx);
   shortcut->setName(shortcutLayerName.c_str());

nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.h

@@ -15,6 +15,6 @@
 
 nvinfer1::ITensor* shortcutLayer(int layerIdx, std::string activation, std::string inputVol, std::string shortcutVol,
     std::map<std::string, std::string>& block, nvinfer1::ITensor* input, nvinfer1::ITensor* shortcut,
-    nvinfer1::INetworkDefinition* network, uint batchSize);
+    nvinfer1::INetworkDefinition* network);
 
 #endif

nvdsinfer_custom_impl_Yolo/layers/slice_layer.cpp

@@ -9,58 +9,72 @@
 
 nvinfer1::ITensor*
 sliceLayer(int layerIdx, std::string& name, nvinfer1::ITensor* input, nvinfer1::Dims start, nvinfer1::Dims size,
-    nvinfer1::Dims stride, nvinfer1::INetworkDefinition* network, uint batchSize)
+    nvinfer1::Dims stride, nvinfer1::INetworkDefinition* network)
 {
   nvinfer1::ITensor* output;
 
-  int tensorBatch = input->getDimensions().d[0];
+  nvinfer1::ISliceLayer* slice;
 
-  nvinfer1::ISliceLayer* slice = network->addSlice(*input, start, size, stride);
+  nvinfer1::Dims inputDims = input->getDimensions();
+
+  if (inputDims.d[0] == -1) {
+    slice = network->addSlice(*input, start, nvinfer1::Dims{}, stride);
+    assert(slice != nullptr);
 
-  if (tensorBatch == -1) {
     int nbDims = size.nbDims;
 
-    nvinfer1::Weights constant1Wt {nvinfer1::DataType::kINT32, nullptr, nbDims};
+    nvinfer1::IShapeLayer* shape = network->addShape(*input);
+    assert(shape != nullptr);
+    std::string shapeLayerName = "shape_" + name + "_" + std::to_string(layerIdx);
+    shape->setName(shapeLayerName.c_str());
+    nvinfer1::ITensor* shapeTensor = shape->getOutput(0);
+    assert(shapeTensor != nullptr);
 
-    int* val1 = new int[nbDims];
-    val1[0] = 1;
-    for (int i = 1; i < nbDims; ++i) {
-      val1[i] = size.d[i];
+#if NV_TENSORRT_MAJOR >= 10
+    nvinfer1::ICastLayer* castShape = network->addCast(*shapeTensor, nvinfer1::DataType::kINT32);
+    assert(castShape != nullptr);
+    std::string castShapeLayerName = "cast_shape_" + name + "_" + std::to_string(layerIdx);
+    castShape->setName(castShapeLayerName.c_str());
+    nvinfer1::ITensor* castShapeTensor = castShape->getOutput(0);
+    assert(castShapeTensor != nullptr);
+    shapeTensor = castShapeTensor;
+#endif
+
+    nvinfer1::Weights constantWt {nvinfer1::DataType::kINT32, nullptr, nbDims};
+
+    int* val = new int[nbDims];
+    for (int i = 0; i < nbDims; ++i) {
+      if (inputDims.d[i] == size.d[i]) {
+        val[i] = 0;
+      }
+      else {
+        val[i] = inputDims.d[i] - size.d[i];
+      }
     }
-    constant1Wt.values = val1;
+    constantWt.values = val;
 
-    nvinfer1::IConstantLayer* constant1 = network->addConstant(nvinfer1::Dims{1, {nbDims}}, constant1Wt);
-    assert(constant1 != nullptr);
-    std::string constant1LayerName = "constant1_" + name + "_" + std::to_string(layerIdx);
-    constant1->setName(constant1LayerName.c_str());
-    nvinfer1::ITensor* constant1Tensor = constant1->getOutput(0);
+    nvinfer1::IConstantLayer* constant = network->addConstant(nvinfer1::Dims{1, {nbDims}}, constantWt);
+    assert(constant != nullptr);
+    std::string constantLayerName = "constant_" + name + "_" + std::to_string(layerIdx);
+    constant->setName(constantLayerName.c_str());
+    nvinfer1::ITensor* constantTensor = constant->getOutput(0);
+    assert(constantTensor != nullptr);
 
-    nvinfer1::Weights constant2Wt {nvinfer1::DataType::kINT32, nullptr, nbDims};
+    nvinfer1::IElementWiseLayer* divide = network->addElementWise(*shapeTensor, *constantTensor,
+        nvinfer1::ElementWiseOperation::kSUB);
+    assert(divide != nullptr);
+    std::string divideLayerName = "divide_" + name + "_" + std::to_string(layerIdx);
+    divide->setName(divideLayerName.c_str());
+    nvinfer1::ITensor* divideTensor = divide->getOutput(0);
+    assert(divideTensor != nullptr);
 
-    int* val2 = new int[nbDims];
-    val2[0] = batchSize;
-    for (int i = 1; i < nbDims; ++i) {
-      val2[i] = 1;
-    }
-    constant2Wt.values = val2;
-
-    nvinfer1::IConstantLayer* constant2 = network->addConstant(nvinfer1::Dims{1, {nbDims}}, constant2Wt);
-    assert(constant2 != nullptr);
-    std::string constant2LayerName = "constant2_" + name + "_" + std::to_string(layerIdx);
-    constant2->setName(constant2LayerName.c_str());
-    nvinfer1::ITensor* constant2Tensor = constant2->getOutput(0);
-
-    nvinfer1::IElementWiseLayer* newSize = network->addElementWise(*constant1Tensor, *constant2Tensor,
-        nvinfer1::ElementWiseOperation::kPROD);
-    assert(newSize != nullptr);
-    std::string newSizeLayerName = "new_size_" + name + "_" + std::to_string(layerIdx);
-    newSize->setName(newSizeLayerName.c_str());
-    nvinfer1::ITensor* newSizeTensor = newSize->getOutput(0);
-
-    slice->setInput(2, *newSizeTensor);
+    slice->setInput(2, *divideTensor);
+  }
+  else {
+    slice = network->addSlice(*input, start, size, stride);
+    assert(slice != nullptr);
   }
 
-  assert(slice != nullptr);
   std::string sliceLayerName = name + "_" + std::to_string(layerIdx);
   slice->setName(sliceLayerName.c_str());
   output = slice->getOutput(0);

nvdsinfer_custom_impl_Yolo/layers/slice_layer.h

@@ -11,6 +11,6 @@
 #include "NvInfer.h"
 
 nvinfer1::ITensor* sliceLayer(int layerIdx, std::string& name, nvinfer1::ITensor* input, nvinfer1::Dims start,
-    nvinfer1::Dims size, nvinfer1::Dims stride, nvinfer1::INetworkDefinition* network, uint batchSize);
+    nvinfer1::Dims size, nvinfer1::Dims stride, nvinfer1::INetworkDefinition* network);
 
 #endif

nvdsinfer_custom_impl_Yolo/layers/upsample_layer.cpp

@@ -24,7 +24,13 @@ upsampleLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1:
   assert(resize != nullptr);
   std::string resizeLayerName = "upsample_" + std::to_string(layerIdx);
   resize->setName(resizeLayerName.c_str());
+  
+#if NV_TENSORRT_MAJOR >= 8 && NV_TENSORRT_MINOR > 0
+  resize->setResizeMode(nvinfer1::InterpolationMode::kNEAREST);
+#else
   resize->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
+#endif
+
   resize->setScales(scale, 4);
   output = resize->getOutput(0);