Add YOLOv8 support

nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.cpp

@@ -3,108 +3,94 @@
  * https://www.github.com/marcoslucianops
  */
 
-#include <math.h>
 #include "batchnorm_layer.h"
 
-nvinfer1::ITensor* batchnormLayer(
-    int layerIdx,
-    std::map<std::string, std::string>& block,
-    std::vector<float>& weights,
-    std::vector<nvinfer1::Weights>& trtWeights,
-    int& weightPtr,
-    std::string weightsType,
-    float eps,
-    nvinfer1::ITensor* input,
+#include <cassert>
+#include <math.h>
+
+nvinfer1::ITensor*
+batchnormLayer(int layerIdx, std::map<std::string, std::string>& block, std::vector<float>& weights,
+    std::vector<nvinfer1::Weights>& trtWeights, int& weightPtr, std::string weightsType, float eps, nvinfer1::ITensor* input,
     nvinfer1::INetworkDefinition* network)
 {
-    nvinfer1::ITensor* output;
+  nvinfer1::ITensor* output;
 
-    assert(block.at("type") == "batchnorm");
-    assert(block.find("filters") != block.end());
+  assert(block.at("type") == "batchnorm");
+  assert(block.find("filters") != block.end());
 
-    int filters = std::stoi(block.at("filters"));
-    std::string activation = block.at("activation");
+  int filters = std::stoi(block.at("filters"));
+  std::string activation = block.at("activation");
 
-    std::vector<float> bnBiases;
-    std::vector<float> bnWeights;
-    std::vector<float> bnRunningMean;
-    std::vector<float> bnRunningVar;
+  std::vector<float> bnBiases;
+  std::vector<float> bnWeights;
+  std::vector<float> bnRunningMean;
+  std::vector<float> bnRunningVar;
 
-    if (weightsType == "weights")
-    {
-        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] + 1.0e-5));
-            weightPtr++;
-        }
+  if (weightsType == "weights") {
+    for (int i = 0; i < filters; ++i) {
+      bnBiases.push_back(weights[weightPtr]);
+      ++weightPtr;
     }
-    else
-    {
-        for (int i = 0; i < filters; ++i)
-        {
-            bnWeights.push_back(weights[weightPtr]);
-            weightPtr++;
-        }
-        for (int i = 0; i < filters; ++i)
-        {
-            bnBiases.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++;
-        }
+    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] + 1.0e-5));
+      ++weightPtr;
+    }
+  }
+  else {
+    for (int i = 0; i < filters; ++i) {
+      bnWeights.push_back(weights[weightPtr]);
+      ++weightPtr;
+    }
+    for (int i = 0; i < filters; ++i) {
+      bnBiases.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;
+    }
+  }
 
-    int 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);
+  int 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(*input, nvinfer1::ScaleMode::kCHANNEL, shift, scale, power);
-    assert(batchnorm != nullptr);
-    std::string batchnormLayerName = "batchnorm_" + std::to_string(layerIdx);
-    batchnorm->setName(batchnormLayerName.c_str());
-    output = batchnorm->getOutput(0);
+  nvinfer1::IScaleLayer* batchnorm = network->addScale(*input, nvinfer1::ScaleMode::kCHANNEL, shift, scale, power);
+  assert(batchnorm != nullptr);
+  std::string batchnormLayerName = "batchnorm_" + std::to_string(layerIdx);
+  batchnorm->setName(batchnormLayerName.c_str());
+  output = batchnorm->getOutput(0);
 
-    output = activationLayer(layerIdx, activation, output, network);
-    assert(output != nullptr);
+  output = activationLayer(layerIdx, activation, output, network);
+  assert(output != nullptr);
 
-    return output;
+  return output;
 }