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

Fix MaxPool layer

Browse Source
This commit is contained in:
Marcos Luciano
2022-07-15 03:57:47 -03:00
parent a19915f44d
commit 0fe1bf6e25
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp
View File

@@ -23,9 +23,11 @@ nvinfer1::ILayer* maxpoolLayer(
assert(pool); assert(pool);
std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx); std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx);
pool->setStrideNd(nvinfer1::Dims{2, {stride, stride}}); pool->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
if (stride == 1) pool->setPaddingNd(nvinfer1::Dims{2, {(size - 1) / 2, (size - 1) / 2}});
if (size == 2 && stride == 1)
{ {
pool->setPaddingNd(nvinfer1::Dims{2, {size / 2, size / 2}}); pool->setPrePadding(nvinfer1::Dims{2, {0, 0}});
pool->setPostPadding(nvinfer1::Dims{2, {1, 1}});
} }
pool->setName(maxpoolLayerName.c_str()); pool->setName(maxpoolLayerName.c_str());