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

Fix ISliceLayer

Browse Source
This commit is contained in:
Marcos Luciano
2023-09-12 00:27:12 -03:00
parent d212d861cb
commit b2388fb3cb
9 changed files with 137 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.cpp
View File

@@ -10,19 +10,20 @@
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)
nvinfer1::INetworkDefinition* network, uint batchSize)
{
nvinfer1::ITensor* output;
assert(block.at("type") == "shortcut");
if (inputVol != shortcutVol) {
nvinfer1::ISliceLayer* slice = network->addSlice(*shortcutInput, nvinfer1::Dims{4, {0, 0, 0, 0}}, input->getDimensions(),
nvinfer1::Dims{4, {1, 1, 1, 1}});
assert(slice != nullptr);
std::string sliceLayerName = "slice_" + std::to_string(layerIdx);
slice->setName(sliceLayerName.c_str());
output = slice->getOutput(0);
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}};
output = sliceLayer(layerIdx, name, shortcutInput, start, size, stride, network, batchSize);
assert(output != nullptr);
}
else
output = shortcutInput;