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

Update upsample_layer.cpp

Browse Source 
replace two matrix multiplication to resize layer, to adapt non square input size.
This commit is contained in:
书包
2021-05-21 08:51:28 +08:00
committed by GitHub
parent 28c6c45ef3
commit 310c139cae
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
non_square/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.cpp
View File

@@ -22,7 +22,7 @@ nvinfer1::ILayer* upsampleLayer(
int w = inpDims.d[2]; int w = inpDims.d[2];
int stride = std::stoi(block.at("stride")); int stride = std::stoi(block.at("stride"));
nvinfer1::Dims preDims{3, /*nvinfer1::Dims preDims{3,
{1, stride * h, h}, {1, stride * h, h},
{nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL, {nvinfer1::DimensionType::kCHANNEL, nvinfer1::DimensionType::kSPATIAL,
nvinfer1::DimensionType::kSPATIAL}}; nvinfer1::DimensionType::kSPATIAL}};
@@ -79,5 +79,12 @@ nvinfer1::ILayer* upsampleLayer(
std::string mm2LayerName = "mm2_" + std::to_string(layerIdx); std::string mm2LayerName = "mm2_" + std::to_string(layerIdx);
mm2->setName(mm2LayerName.c_str()); mm2->setName(mm2LayerName.c_str());
return mm2; return mm2;*/
nvinfer1::IResizeLayer* resize_layer = network->addResize(*input);
resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
float scale[3] = {1, stride, stride};
resize_layer->setScales(scale, 3);
std::string layer_name = "upsample_" + std::to_string(layerIdx);
resize_layer->setName(layer_name.c_str());
return resize_layer;
} }