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

pose int8 images path and fix pose p6 (#1507)

Browse Source 
* pose int8 images path and fix pose p6

* pose int8 images path and fix pose p6
This commit is contained in:
lindsayshuo 2024-04-29 14:30:27 +08:00 committed by GitHub
parent ae38073988
commit 37633b4529
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
yolov8/src/model.cpp
View File

@ -1762,7 +1762,7 @@ nvinfer1::IHostMemory* buildEngineYolov8PoseP6(nvinfer1::IBuilder* builder, nvin
(kInputH / strides[0]) * (kInputW / strides[0]), gw, "pose"); (kInputH / strides[0]) * (kInputW / strides[0]), gw, "pose");
nvinfer1::ITensor* inputTensor30_dfl_0[] = {dfl30_0->getOutput(0), split30_0_1->getOutput(0), nvinfer1::ITensor* inputTensor30_dfl_0[] = {dfl30_0->getOutput(0), split30_0_1->getOutput(0),
shuffle_conv20->getOutput(0)}; shuffle_conv20->getOutput(0)};
nvinfer1::IConcatenationLayer* cat30_dfl_0 = network->addConcatenation(inputTensor30_dfl_0, 2); nvinfer1::IConcatenationLayer* cat30_dfl_0 = network->addConcatenation(inputTensor30_dfl_0, 3);
// P4 processing steps (remains unchanged) // P4 processing steps (remains unchanged)
nvinfer1::IShuffleLayer* shuffle30_1 = nvinfer1::IShuffleLayer* shuffle30_1 =
@ -1783,7 +1783,7 @@ nvinfer1::IHostMemory* buildEngineYolov8PoseP6(nvinfer1::IBuilder* builder, nvin
(kInputH / strides[1]) * (kInputW / strides[1]), gw, "pose"); (kInputH / strides[1]) * (kInputW / strides[1]), gw, "pose");
nvinfer1::ITensor* inputTensor30_dfl_1[] = {dfl30_1->getOutput(0), split30_1_1->getOutput(0), nvinfer1::ITensor* inputTensor30_dfl_1[] = {dfl30_1->getOutput(0), split30_1_1->getOutput(0),
shuffle_conv23->getOutput(0)}; shuffle_conv23->getOutput(0)};
nvinfer1::IConcatenationLayer* cat30_dfl_1 = network->addConcatenation(inputTensor30_dfl_1, 2); nvinfer1::IConcatenationLayer* cat30_dfl_1 = network->addConcatenation(inputTensor30_dfl_1, 3);
// P5 processing steps (remains unchanged) // P5 processing steps (remains unchanged)
nvinfer1::IShuffleLayer* shuffle30_2 = nvinfer1::IShuffleLayer* shuffle30_2 =
@ -1804,7 +1804,7 @@ nvinfer1::IHostMemory* buildEngineYolov8PoseP6(nvinfer1::IBuilder* builder, nvin
(kInputH / strides[2]) * (kInputW / strides[2]), gw, "pose"); (kInputH / strides[2]) * (kInputW / strides[2]), gw, "pose");
nvinfer1::ITensor* inputTensor30_dfl_2[] = {dfl30_2->getOutput(0), split30_2_1->getOutput(0), nvinfer1::ITensor* inputTensor30_dfl_2[] = {dfl30_2->getOutput(0), split30_2_1->getOutput(0),
shuffle_conv26->getOutput(0)}; shuffle_conv26->getOutput(0)};
nvinfer1::IConcatenationLayer* cat30_dfl_2 = network->addConcatenation(inputTensor30_dfl_2, 2); nvinfer1::IConcatenationLayer* cat30_dfl_2 = network->addConcatenation(inputTensor30_dfl_2, 3);
// P6 processing steps // P6 processing steps
nvinfer1::IShuffleLayer* shuffle30_3 = network->addShuffle(*cat30_3->getOutput(0)); nvinfer1::IShuffleLayer* shuffle30_3 = network->addShuffle(*cat30_3->getOutput(0));
@ -1819,16 +1819,16 @@ nvinfer1::IHostMemory* buildEngineYolov8PoseP6(nvinfer1::IBuilder* builder, nvin
DFL(network, weightMap, *split30_3_0->getOutput(0), 4, (kInputH / strides[3]) * (kInputW / strides[3]), 1, DFL(network, weightMap, *split30_3_0->getOutput(0), 4, (kInputH / strides[3]) * (kInputW / strides[3]), 1,
1, 0, "model.30.dfl.conv.weight"); 1, 0, "model.30.dfl.conv.weight");
// det2 // det3
auto shuffle_conv29 = cv4_conv_combined(network, weightMap, *conv29->getOutput(0), "model.30.cv4.3", auto shuffle_conv29 = cv4_conv_combined(network, weightMap, *conv29->getOutput(0), "model.30.cv4.3",
(kInputH / strides[3]) * (kInputW / strides[3]), gw, "pose"); (kInputH / strides[3]) * (kInputW / strides[3]), gw, "pose");
nvinfer1::ITensor* inputTensor30_dfl_3[] = {dfl30_3->getOutput(0), split30_3_1->getOutput(0), nvinfer1::ITensor* inputTensor30_dfl_3[] = {dfl30_3->getOutput(0), split30_3_1->getOutput(0),
shuffle_conv29->getOutput(0)}; shuffle_conv29->getOutput(0)};
nvinfer1::IConcatenationLayer* cat30_dfl_3 = network->addConcatenation(inputTensor30_dfl_3, 2); nvinfer1::IConcatenationLayer* cat30_dfl_3 = network->addConcatenation(inputTensor30_dfl_3, 3);
nvinfer1::IPluginV2Layer* yolo = addYoLoLayer( nvinfer1::IPluginV2Layer* yolo = addYoLoLayer(
network, std::vector<nvinfer1::IConcatenationLayer*>{cat30_dfl_0, cat30_dfl_1, cat30_dfl_2, cat30_dfl_3}, network, std::vector<nvinfer1::IConcatenationLayer*>{cat30_dfl_0, cat30_dfl_1, cat30_dfl_2, cat30_dfl_3},
strides, stridesLength, false, false); strides, stridesLength, false, true);
yolo->getOutput(0)->setName(kOutputTensorName); yolo->getOutput(0)->setName(kOutputTensorName);
network->markOutput(*yolo->getOutput(0)); network->markOutput(*yolo->getOutput(0));
@ -1841,8 +1841,8 @@ nvinfer1::IHostMemory* buildEngineYolov8PoseP6(nvinfer1::IBuilder* builder, nvin
std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
assert(builder->platformHasFastInt8()); assert(builder->platformHasFastInt8());
config->setFlag(nvinfer1::BuilderFlag::kINT8); config->setFlag(nvinfer1::BuilderFlag::kINT8);
auto* calibrator = auto* calibrator = new Int8EntropyCalibrator2(1, kInputW, kInputH, kInputQuantizationFolder, "int8calib.table",
new Int8EntropyCalibrator2(1, kInputW, kInputH, "../coco_calib/", "int8calib.table", kInputTensorName); kInputTensorName);
config->setInt8Calibrator(calibrator); config->setInt8Calibrator(calibrator);
#endif #endif