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

New YOLOv5 conversion and support (>= v2.0)

Browse Source
This commit is contained in:
Marcos Luciano
2022-07-14 13:58:55 -03:00
parent ea31a54bbb
commit 8983e7e99e
2 changed files with 4 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp
View File

@@ -18,18 +18,14 @@ nvinfer1::ILayer* maxpoolLayer(
int size = std::stoi(block.at("size")); int size = std::stoi(block.at("size"));
int stride = std::stoi(block.at("stride")); int stride = std::stoi(block.at("stride"));
int pad = 0;
if (block.find("pad") != block.end())
{
pad = std::stoi(block.at("pad"));
}
nvinfer1::IPoolingLayer* pool nvinfer1::IPoolingLayer* pool
= network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX, nvinfer1::Dims{2, {size, size}}); = network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX, nvinfer1::Dims{2, {size, size}});
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}});
pool->setPaddingNd(nvinfer1::Dims{2, {pad, pad}}); if (stride == 1) {
pool->setPaddingNd(nvinfer1::Dims{2, {size / 2, size / 2}});
}
pool->setName(maxpoolLayerName.c_str()); pool->setName(maxpoolLayerName.c_str());
return pool; return pool;

5
utils/gen_wts_yoloV5.py
View File

@@ -210,16 +210,13 @@ class Layers(object):
stride = m.stride stride = m.stride
size = m.kernel_size size = m.kernel_size
pad = m.padding
mode = m.ceil_mode mode = m.ceil_mode
m = 'maxpool_up' if mode else 'maxpool' m = 'maxpool_up' if mode else 'maxpool'
p = 'pad=%d\n' % pad if pad > 0 else ''
self.fc.write('\n[%s]\n' % m + self.fc.write('\n[%s]\n' % m +
'stride=%d\n' % stride + 'stride=%d\n' % stride +
'size=%d\n' % size + 'size=%d\n' % size)
p)
def upsample(self, child): def upsample(self, child):
self.blocks[self.current] += 1 self.blocks[self.current] += 1