haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix strides (#1492)

Browse Source 
* yolov8 p2

* yolov8 p2

* yolov8 p2

* yolov8 p2

* Update yolov8_det.cpp

* Update model.cpp

* Update model.cpp

* Update model.cpp

* Update model.cpp

* fix strides
This commit is contained in:
lindsayshuo 2024-04-23 10:22:29 +08:00 committed by GitHub
parent d1a184ffca
commit d033a63f07
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 21 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
yolov8/src/model.cpp
View File

@ -20,6 +20,15 @@ static int get_depth(int x, float gd) {
return std::max<int>(r, 1);
}
void calculateStrides(nvinfer1::IElementWiseLayer* conv_layers[], int size, int reference_size, int strides[]) {
for (int i = 0; i < size; ++i) {
nvinfer1::ILayer* layer = conv_layers[i];
nvinfer1::Dims dims = layer->getOutput(0)->getDimensions();
int feature_map_size = dims.d[1];
strides[i] = reference_size / feature_map_size;
}
}
static nvinfer1::IElementWiseLayer* Proto(nvinfer1::INetworkDefinition* network,
std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input,
std::string lname, float gw, int max_channels) {
@ -220,7 +229,9 @@ nvinfer1::IHostMemory* buildEngineYolov8Det(nvinfer1::IBuilder* builder, nvinfer
********************************************* YOLOV8 DETECT ******************************************
*******************************************************************************************************/
int strides[] = {8, 16, 32};
nvinfer1::IElementWiseLayer* conv_layers[] = {conv3, conv5, conv7};
int strides[sizeof(conv_layers) / sizeof(conv_layers[0])];
calculateStrides(conv_layers, sizeof(conv_layers) / sizeof(conv_layers[0]), kInputH, strides);
int stridesLength = sizeof(strides) / sizeof(int);
nvinfer1::IShuffleLayer* shuffle22_0 = network->addShuffle(*cat22_0->getOutput(0));
@ -507,7 +518,9 @@ nvinfer1::IHostMemory* buildEngineYolov8DetP6(nvinfer1::IBuilder* builder, nvinf
/*******************************************************************************************************
********************************************* YOLOV8 DETECT ******************************************
*******************************************************************************************************/
int strides[] = {8, 16, 32, 64};
nvinfer1::IElementWiseLayer* conv_layers[] = {conv3, conv5, conv7, conv9};
int strides[sizeof(conv_layers) / sizeof(conv_layers[0])];
calculateStrides(conv_layers, sizeof(conv_layers) / sizeof(conv_layers[0]), kInputH, strides);
int stridesLength = sizeof(strides) / sizeof(int);
// P3 processing steps (remains unchanged)
@ -817,7 +830,9 @@ nvinfer1::IHostMemory* buildEngineYolov8DetP2(nvinfer1::IBuilder* builder, nvinf
********************************************* YOLOV8 DETECT ******************************************
*******************************************************************************************************/
int strides[] = {4, 8, 16, 32};
nvinfer1::IElementWiseLayer* conv_layers[] = {conv1, conv3, conv5, conv7};
int strides[sizeof(conv_layers) / sizeof(conv_layers[0])];
calculateStrides(conv_layers, sizeof(conv_layers) / sizeof(conv_layers[0]), kInputH, strides);
int stridesLength = sizeof(strides) / sizeof(int);
// P2 processing steps (remains unchanged)
@ -1148,7 +1163,9 @@ nvinfer1::IHostMemory* buildEngineYolov8Seg(nvinfer1::IBuilder* builder, nvinfer
********************************************* YOLOV8 DETECT ******************************************
*******************************************************************************************************/
int strides[] = {8, 16, 32};
nvinfer1::IElementWiseLayer* conv_layers[] = {conv3, conv5, conv7};
int strides[sizeof(conv_layers) / sizeof(conv_layers[0])];
calculateStrides(conv_layers, sizeof(conv_layers) / sizeof(conv_layers[0]), kInputH, strides);
int stridesLength = sizeof(strides) / sizeof(int);
nvinfer1::IShuffleLayer* shuffle22_0 = network->addShuffle(*cat22_0->getOutput(0));