thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add PP-YOLOE support

Browse Source
This commit is contained in:
Marcos Luciano
2022-07-24 18:00:47 -03:00
parent d09879d557
commit a3782ed65e
51 changed files with 1812 additions and 600 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
nvdsinfer_custom_impl_Yolo/layers/implicit_layer.cpp
View File

@@ -5,17 +5,25 @@
#include "implicit_layer.h"
nvinfer1::ILayer* implicitLayer(
int channels,
nvinfer1::ITensor* implicitLayer(
int layerIdx,
std::map<std::string, std::string>& block,
std::vector<float>& weights,
std::vector<nvinfer1::Weights>& trtWeights,
int& weightPtr,
nvinfer1::INetworkDefinition* network)
{
nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, channels};
nvinfer1::ITensor* output;
float* val = new float[channels];
for (int i = 0; i < channels; ++i)
assert(block.at("type") == "implicit_add" || block.at("type") == "implicit_mul");
assert(block.find("filters") != block.end());
int filters = std::stoi(block.at("filters"));
nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, filters};
float* val = new float[filters];
for (int i = 0; i < filters; ++i)
{
val[i] = weights[weightPtr];
weightPtr++;
@@ -23,8 +31,11 @@ nvinfer1::ILayer* implicitLayer(
convWt.values = val;
trtWeights.push_back(convWt);
nvinfer1::IConstantLayer* implicit = network->addConstant(nvinfer1::Dims3{static_cast<int>(channels), 1, 1}, convWt);
nvinfer1::IConstantLayer* implicit = network->addConstant(nvinfer1::Dims{3, {filters, 1, 1}}, convWt);
assert(implicit != nullptr);
std::string implicitLayerName = block.at("type") + "_" + std::to_string(layerIdx);
implicit->setName(implicitLayerName.c_str());
output = implicit->getOutput(0);
return implicit;
return output;
}