TensorRt（6）yolov3.weight转换、onnx_graphsurgeon和c++ api实现添加NMS

news2024/11/25 4:46:26

前面博文 【opencv dnn模块 示例(3) 目标检测 object_detection (2) YOLO object detection】 介绍了 使用opencv dnn模块加载yolo weights格式模型的详细说明。 又在博文 【TensorRt（4）yolov3加载测试】 说明了如何将onnx编译为tensorrt格式并使用的方式，但是存在后处理繁琐的问题。本文将继续说明tensorrt加载yolov3模型进行推理，并且增加efficient nms trt模块。

文章目录

  • 1、yolov3.weight模型格式转换
  • 2、添加 EfficientNMS_TRT 层
    • 2.1、添加 Reshape层
    • 2.2、添加 EfficentNMS 层
    • 2.3、使用 trtexec 导出优化后的模型
    • 2.4、Tensorrt API c++测试 yolo_nms
  • 3、c++ API实现添加NMS并推理测试
    • 3.1、扩展维度
    • 3.2、添加nms层
    • 3.3、编译、测试代码

1、yolov3.weight模型格式转换

我们先下载项目 Yolov3_Darknet_PyTorch_Onnx_Converter，配置好环境周将weights格式进行转换onnx格式，以yolov3.weights为例

python converter.py yolov3.cfg yolov3.weights 416 416

之后，会保存pt和onnx两个转换后的模型文件。通过netron查看yolov3.onnx信息，常规的三个分支输出被整合，并且仅保留80个classes 信息和 4个rect信息。

在这里插入图片描述
可以直接使用yolov3博文中的opencv dnn 测试代码，仅将后处理修改为如下即可

void postprocess(Mat& frame, const std::vector<Mat>& outs, Net& net)
{
    std::vector<int> classIds;
    std::vector<float> confidences;
    std::vector<Rect> boxes;
    
	for(int i = 0; i < outs[0].rows; i++)  
	{
        Mat scores = outs[1].row(i);
        Point classIdPoint;
        double confidence;
        minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);
        
        if(confidence > confThreshold) {
            Mat rect = outs[0].row(i);
            int centerX = (int)(rect.at<float>(0) * frame.cols);
            int centerY = (int)(rect.at<float>(1) * frame.rows);
            int width = (int)(rect.at<float>(2) * frame.cols);
            int height = (int)(rect.at<float>(3) * frame.rows);
            int left = centerX - width / 2;
            int top = centerY - height / 2;

            classIds.push_back(classIdPoint.x);
            confidences.push_back((float)confidence);
            boxes.push_back(Rect(left, top, width, height));
        }
    }

    std::vector<int> indices;
    NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, indices);
    for(size_t i = 0; i < indices.size(); ++i) {
        int idx = indices[i];
        Rect box = boxes[idx];
        drawPred(classIds[idx], confidences[idx], box.x, box.y,
                 box.x + box.width, box.y + box.height, frame);
    }
}

测试结果正常

2、添加 EfficientNMS_TRT 层

在前面生成的 yolov3.onnx 模型后，通过 onnx_graphsurgeon 为其添加层类型为 EfficientNMS_TRT 的 nms处理模块。

注意： onnx模型添加EfficientNMS_TRT 仅支持 Tensorrt 作为后端推理。

首先我们需要知道EfficientNMS_TRT的输入和输出格式（更多信息点击查看），

  • 输入：
    • boxes : [batch_size, number_boxes, 4] or [batch_size, number_boxes, number_classes, 4]
    • scores : [batch_size, number_boxes, number_classes]
  • 输出：
    • num_detections: [batch_size, 1], int32
    • detection_boxes: [batch_size, max_output_boxes, 4]
    • detection_scores: [batch_size, max_output_boxes]
    • detection_classes: [batch_size, max_output_boxes], int32

2.1、添加 Reshape层

yolov3.onnx 的输出是 10647x4 和 10647x80，分别来自节点 /Concat_3 和 /Concat 而作为EfficientNMS_TRT输入仅缺少一个维度用来描述batch_size，即这里要求输入为 1x10647x4 和 1x10647x80。

我们使用 onnx_graphsurgeon 库先为onnx添加2个 Reshape节点以扩展原来两个输出的维度。

import onnx_graphsurgeon as gs
import numpy as np
import onnx

def get_nms_input(graph, class_num=1, output_name="Concat_"):
    # # input
    origin_outputs_node = [node for node in graph.nodes if node.name in output_name]
    
    #num_dections = origin_outputs_node[0].outputs[0].shape[0]
    #num_classes = origin_outputs_node[1].outputs[0].shape[1]

    var_classes = gs.Variable(name="classes",shape=[1] + origin_outputs_node[0].outputs[0].shape, dtype=np.float32)
    var_boxes = gs.Variable(name="boxes",shape=[1] + origin_outputs_node[1].outputs[0].shape, dtype=np.float32)

    origin_outputs_node[0].outputs[0].name = "classes_ori"  # 不能同名
    origin_outputs_node[1].outputs[0].name = "boxes_ori"
    
    node_classes = gs.Node(op='Reshape', name='reshape_classes',
                            inputs=[origin_outputs_node[0].outputs[0], 
                                    gs.Constant('new_classes_shape', values=np.array(var_classes.shape, dtype=np.int64))],
                            outputs=[var_classes])
    node_boxes = gs.Node(op='Reshape', name='reshape_boxes',
                            inputs=[origin_outputs_node[1].outputs[0], 
                                    gs.Constant('new_boxes_shape', values=np.array(var_boxes.shape, dtype=np.int64))],
                           outputs=[var_boxes])

    graph.nodes.extend([node_classes, node_boxes,])
    graph.outputs = [node_classes.outputs[0], node_boxes.outputs[0],]

    graph.cleanup().toposort()
    onnx.save(gs.export_onnx(graph),"./yolov3_extend.onnx")

	## 中间结果保存
    # graph.cleanup().toposort()
    # onnx.save(gs.export_onnx(graph),"./yolov3_extend.onnx")

    return graph

if __name__ == "__main__":
    onnx_path = "yolov3.onnx"
    graph = gs.import_onnx(onnx.load(onnx_path))

    # # 添加op得到Efficient NMS plugin的input
    graph = get_nms_input(graph,class_num=80,output_name=("/Concat_3","/Concat"))

我们运行后查看中间结果
在这里插入图片描述
使用opencv dnn 或 使用tensorrt 导出优化编译模型均正常。

2.2、添加 EfficentNMS 层

在2.1中的代码下，添加一个函数用于在main函数中为onnx模型增加nms功能。

def create_and_add_plugin_node(graph, max_output_boxes):
    
    batch_size = graph.inputs[0].shape[0]
    print("The batch size is: ", batch_size)
    input_h = graph.inputs[0].shape[2]
    input_w = graph.inputs[0].shape[3]

    tensors = graph.tensors()
    boxes_tensor = tensors["boxes"]  
    confs_tensor = tensors["classes"]

    num_detections = gs.Variable(name="num_detections").to_variable(dtype=np.int32, shape=[batch_size, 1])
    nmsed_boxes = gs.Variable(name="detection_boxes").to_variable(dtype=np.float32, shape=[batch_size, max_output_boxes, 4])
    nmsed_scores = gs.Variable(name="detection_scores").to_variable(dtype=np.float32, shape=[batch_size, max_output_boxes])
    nmsed_classes = gs.Variable(name="detection_classes").to_variable(dtype=np.int32, shape=[batch_size, max_output_boxes])

    new_outputs = [num_detections, nmsed_boxes, nmsed_scores, nmsed_classes]

    mns_node = gs.Node(
        op="EfficientNMS_TRT",
        attrs=create_attrs(max_output_boxes),
        inputs=[boxes_tensor, confs_tensor],
        outputs=new_outputs)

    graph.nodes.append(mns_node)
    graph.outputs = new_outputs

    return graph.cleanup().toposort()


def create_attrs(max_output_boxes=100):

    attrs = {}

    attrs["score_threshold"] = 0.25 
    attrs["iou_threshold"] = 0.45  
    attrs["max_output_boxes"] = max_output_boxes
    attrs["background_class"] = -1
    attrs["score_activation"] = False
    attrs["class_agnostic"] = True
    attrs["box_coding"] = 1
    # 001 is the default plugin version the parser will search for, and therefore can be omitted,
    # but we include it here for illustrative purposes.
    attrs["plugin_version"] = "1"

    return attrs


if __name__ == "__main__":
    onnx_path = "yolov3.onnx"
    graph = gs.import_onnx(onnx.load(onnx_path))

    # # 添加op得到Efficient NMS plugin的input
    graph = get_nms_input(graph,class_num=80,output_name=("/Concat_3","/Concat"))

    # # 添加Efficient NMS plugin
    graph = create_and_add_plugin_node(graph, 100)

    # # 保存图结构
    onnx.save(gs.export_onnx(graph),"yolov3_nms.onnx")

重新运行程序，查看添加后onnx网络结构如下：
在这里插入图片描述

2.3、使用 trtexec 导出优化后的模型

使用命令如下：

trtexec.exe --onnx=E:\DeepLearning\darknet-yolo3-master\build\darknet\x64\yolov3_nms.onnx --saveEngine=E:\DeepLearning\darknet-yolo3-master\build\darknet\x64\yolov3_nms.onnx.engine --fp16

转换成功信息如下，可以明确看到有1个输入，4个输出。

[06/25/2024-12:25:46] [I] Using random values for input images
[06/25/2024-12:25:46] [I] Input binding for images with dimensions 1x3x416x416 is created.
[06/25/2024-12:25:46] [I] Output binding for num_detections with dimensions 1x1 is created.
[06/25/2024-12:25:46] [I] Output binding for detection_boxes with dimensions 1x100x4 is created.
[06/25/2024-12:25:46] [I] Output binding for detection_scores with dimensions 1x100 is created.
[06/25/2024-12:25:46] [I] Output binding for detection_classes with dimensions 1x100 is created.
[06/25/2024-12:25:46] [I] Starting inference
[06/25/2024-12:25:49] [I] Warmup completed 23 queries over 200 ms
[06/25/2024-12:25:49] [I] Timing trace has 347 queries over 3.01801 s
[06/25/2024-12:25:49] [I]
[06/25/2024-12:25:49] [I] === Trace details ===
[06/25/2024-12:25:49] [I] Trace averages of 10 runs:
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.5421 ms - Host latency: 8.74631 ms (enqueue 0.733459 ms)
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.65691 ms - Host latency: 8.86809 ms (enqueue 0.754318 ms)
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.69078 ms - Host latency: 8.89894 ms (enqueue 0.720746 ms)
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.53977 ms - Host latency: 8.74949 ms (enqueue 0.887427 ms)
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.71277 ms - Host latency: 8.92263 ms (enqueue 0.95188 ms)
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.58589 ms - Host latency: 8.79873 ms (enqueue 0.861426 ms)
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.5209 ms - Host latency: 8.7323 ms (enqueue 0.879272 ms)
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.62339 ms - Host latency: 8.83303 ms (enqueue 0.821729 ms)
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.66113 ms - Host latency: 8.86626 ms (enqueue 0.718994 ms)
[06/25/2024-12:25:49] [I] Average on 10 runs - GPU latency: 8.53975 ms - Host latency: 8.74404 ms (enqueue 0.683179 ms)
[06/25/2024-12:25:49] [I]
[06/25/2024-12:25:49] [I] === Performance summary ===
[06/25/2024-12:25:49] [I] Throughput: 114.976 qps
[06/25/2024-12:25:49] [I] Latency: min = 8.62891 ms, max = 9.82397 ms, mean = 8.8045 ms, median = 8.75854 ms, percentile(90%) = 8.92554 ms, percentile(95%) = 9.17676 ms, percentile(99%) = 9.52246 ms
[06/25/2024-12:25:49] [I] Enqueue Time: min = 0.532227 ms, max = 2.55029 ms, mean = 0.750556 ms, median = 0.699585 ms, percentile(90%) = 0.886719 ms, percentile(95%) = 1.04114 ms, percentile(99%) = 1.60669 ms
[06/25/2024-12:25:49] [I] H2D Latency: min = 0.176514 ms, max = 0.218506 ms, mean = 0.190991 ms, median = 0.190277 ms, percentile(90%) = 0.199463 ms, percentile(95%) = 0.20459 ms, percentile(99%) = 0.213135 ms
[06/25/2024-12:25:49] [I] GPU Compute Time: min = 8.41919 ms, max = 9.6019 ms, mean = 8.59789 ms, median = 8.55249 ms, percentile(90%) = 8.70782 ms, percentile(95%) = 8.96716 ms, percentile(99%) = 9.30908 ms
[06/25/2024-12:25:49] [I] D2H Latency: min = 0.00976563 ms, max = 0.0393066 ms, mean = 0.0156149 ms, median = 0.015625 ms, percentile(90%) = 0.0187988 ms, percentile(95%) = 0.019165 ms, percentile(99%) = 0.0251465 ms
[06/25/2024-12:25:49] [I] Total Host Walltime: 3.01801 s
[06/25/2024-12:25:49] [I] Total GPU Compute Time: 2.98347 s
[06/25/2024-12:25:49] [W] * GPU compute time is unstable, with coefficient of variance = 1.99113%.
[06/25/2024-12:25:49] [W]   If not already in use, locking GPU clock frequency or adding --useSpinWait may improve the stability.
[06/25/2024-12:25:49] [I] Explanations of the performance metrics are printed in the verbose logs.
[06/25/2024-12:25:49] [I]
&&&& PASSED TensorRT.trtexec [TensorRT v8601] # trtexec.exe --onnx=E:\DeepLearning\darknet-yolo3-master\build\darknet\x64\yolov3_nms.onnx --saveEngine=E:\DeepLearning\darknet-yolo3-master\build\darknet\x64\yolov3_nms.onnx.engine --fp16

2.4、Tensorrt API c++测试 yolo_nms

由于EfficientNMS_TRT 层属于 tensorrt的plugins中的实现，因此在解析模型文件时，需要调用nvinfer1:initLibNvInferPlugins()函数的加载开启插件模块，前向测试代码如下：

int inference_nms()
{
    Logger logger(nvinfer1::ILogger::Severity::kVERBOSE);

    std::string trtFile = R"(E:\DeepLearning\darknet-yolo3-master\build\darknet\x64\yolov3_nms.onnx.engine)";

    std::ifstream ifs(trtFile, std::ifstream::binary);
    if(!ifs) {
        return false;
    }

    ifs.seekg(0, std::ios_base::end);
    int size = ifs.tellg();
    ifs.seekg(0, std::ios_base::beg);

    std::unique_ptr<char> pData(new char[size]);
    ifs.read(pData.get(), size);
    ifs.close();

    // engine模型
    nvinfer1:initLibNvInferPlugins(&logger, ""); // 加载插件模块

    std::shared_ptr<nvinfer1::ICudaEngine> mEngine;

    SampleUniquePtr<nvinfer1::IRuntime> runtime{nvinfer1::createInferRuntime(logger.getTRTLogger())};
    mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(
        runtime->deserializeCudaEngine(pData.get(), size), InferDeleter());

    auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());

    std::vector<void*> bindings(mEngine->getNbBindings());
    for(int i = 0; i < bindings.size(); i++) {
        nvinfer1::DataType type = mEngine->getBindingDataType(i);
        nvinfer1::Dims dims = mEngine->getBindingDimensions(i);
        size_t volume = std::accumulate(dims.d, dims.d + dims.nbDims, 1, std::multiplies<size_t>());
        switch(type) {
            case nvinfer1::DataType::kINT32:
            case nvinfer1::DataType::kFLOAT: volume *= 4; break;  // 明确为类型 float
            case nvinfer1::DataType::kHALF: volume *= 2; break;
            case nvinfer1::DataType::kBOOL:
            case nvinfer1::DataType::kINT8:
            default:break;
        }
        CHECK(cudaMalloc(&bindings[i], volume));
    }

    // 输入
    cv::Mat img = cv::imread(R"(E:\DeepLearning\yolov5\data\images\bus.jpg)");
    cv::Mat blob = cv::dnn::blobFromImage(img, 1 / 255., cv::Size(inpWidth, inpHeight), {0,0,0}, true, false);

    // 推理
    CHECK(cudaMemcpy(bindings[0], static_cast<const void*>(blob.data), 1 * 3 * 416 * 416 * sizeof(float), cudaMemcpyHostToDevice));
    context->executeV2(bindings.data());
    context->executeV2(bindings.data());
    context->executeV2(bindings.data());
    context->executeV2(bindings.data());

    auto t1 = cv::getTickCount();
    
    CHECK(cudaMemcpy(bindings[0], static_cast<const void*>(blob.data), 1 * 3 * 416 * 416 * sizeof(float), cudaMemcpyHostToDevice));
    context->executeV2(bindings.data());
    float h_input[inpWidth * inpHeight * 3];
    int h_output_0[1];   //1
    float h_output_1[1 * 100 * 4];   //1
    float h_output_2[1 * 100];   //1
    int h_output_3[1 * 100];   //1

    cudaMemcpy(h_output_0, bindings[1], 1 * sizeof(int), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_output_1, bindings[2], 1 * 100 * 4 * sizeof(float), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_output_2, bindings[3], 1 * 100 * sizeof(float), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_output_3, bindings[4], 1 * 100 * sizeof(int), cudaMemcpyDeviceToHost);

    auto t2 = cv::getTickCount();
    std::string label = cv::format("inference time: %.2f ms", (t2 - t1) / cv::getTickFrequency() * 1000);
    std::cout << label << std::endl;
    cv::putText(img, label, cv::Point(10, 50), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0));

    // 后处理
    float x_factor = (float)img.cols / inpWidth;
    float y_factor = (float)img.rows / inpHeight;

    int numbers = *h_output_0;
    for(int i = 0; i < numbers; i++) {

        float conf = h_output_2[i];
        if(conf < scoreThreshold)
            continue;

        int clsId = h_output_3[i];

        float x = h_output_1[i * 4];
        float y = h_output_1[i * 4 + 1];
        float w = h_output_1[i * 4 + 2];
        float h = h_output_1[i * 4 + 3];

        drawPred(clsId, conf, x * 416 * x_factor, y * 416 * y_factor, w * 416 * x_factor, h * 416 * y_factor, img);
    }
    cv::imshow("res", img);
    cv::waitKey();

    // 资源释放
    cudaFree(bindings[0]);
    cudaFree(bindings[1]);

    return 0;
}

3、c++ API实现添加NMS并推理测试

直接从yolov3.onnx模型开始，流程与上面一样，这里使用 Shuffle 来进行扩展维度，再添加EfficientNMS_TRT。由于使用c++接口，使用变复杂。

若从yolov3_extend.onnx模型开始，那么仅需要添加nms层即可（注意 输入输出层的标记）。

3.1、扩展维度

这里不解释，在某个tensor后添加一个reshape层，再将这个tensorrt的输出标记为非输出节点（已确定不再是输出），并返回节点的输出tensor。

nvinfer1::ITensor* unsqueezeTensor(SampleUniquePtr<nvinfer1::INetworkDefinition>& network, nvinfer1::ITensor* inputTensor, int axis)
{
    // 获取输入张量的维度信息
    nvinfer1::Dims dims = inputTensor->getDimensions();
    // 在指定的轴上插入一个维度
    for(int i = dims.nbDims; i > axis; --i) {
        dims.d[i] = dims.d[i - 1];
    }
    dims.d[axis] = 1;
    dims.nbDims++;

    // 重新配置张量的维度
    auto layer = network->addShuffle(*inputTensor);
    layer->setReshapeDimensions(dims);
    return layer->getOutput(0);
}

3.2、添加nms层

python中通过onnx_graphsurgeon 添加EfficientNMS_TRT 非常简单，设定参数，明确输出输出即可。c++中需要如下几个步骤:

    1. 加载插件，获取插件注册器 IPluginRegistry
    1. 通过创建器，获取指定名称( “EfficientNMS_TRT”) 插件的创建器 IPluginCreator
    1. 使用 PluginField 结构指定插件的参数数组
    1. 准备 创建器 需要的参数 PluginFieldCollection，其由 PluginField 数组构造
    1. 创建一个插件对象 IPlugin
    1. 添加插件对象到 network 中
    1. 设置添加的层的输出的名称，并标记为输出

在第三步中，PluginField 的参数类型如何设定，需要参考 对应层EfficientNMS_TRT的源代码，可以参考源代码efficientNMSPlugin.cpp 的参数说明和创建代码

在这里插入图片描述
本文中添加插件层的代码实现，如下

    // 1
	nvinfer1:initLibNvInferPlugins(&logger, "");
    auto registry = nvinfer1::getBuilderPluginRegistry(nvinfer1::EngineCapability::kDEFAULT);
    // 2
    auto creator = registry->getPluginCreator("EfficientNMS_TRT", "1");  	
	// 3
    auto score_threshold = 0.25f;
    auto iou_threshold = 0.45f;
    auto max_output_boxes = 100;
    auto background_class = -1;
    auto score_activation = 0;
    auto class_agnostic = 1;
    auto box_coding = 1;
    //auto plugin_version = "1";       
    std::vector<nvinfer1::PluginField> new_pluginData_list = {
        {"score_threshold", &score_threshold, PluginFieldType::kFLOAT32, sizeof(score_threshold)},
        {"iou_threshold", &iou_threshold, PluginFieldType::kFLOAT32, sizeof(iou_threshold)},
        {"max_output_boxes", &max_output_boxes, PluginFieldType::kINT32, sizeof(max_output_boxes)},
        {"background_class", &background_class, PluginFieldType::kINT32, sizeof(background_class)},
        {"score_activation", &score_activation, PluginFieldType::kINT32, sizeof(score_activation)},
        {"class_agnostic", &class_agnostic, PluginFieldType::kINT32, sizeof(class_agnostic)},
        {"box_coding", &box_coding, PluginFieldType::kINT32, sizeof(box_coding)},
        //{"plugin_version", &plugin_version, PluginFieldType::kCHAR, sizeof(plugin_version)},
    };	
	// 4
    auto collection = nvinfer1::PluginFieldCollection();
    collection.nbFields = new_pluginData_list.size();
    collection.fields = new_pluginData_list.data();
	// 5
    auto plugin = creator->createPlugin("nms_layer", &collection);
	// 6
    auto layer = network->addPluginV2(newOutputsTensors.data(), newOutputsTensors.size(), *plugin);
    layer->getOutput(0)->setName("det_num");
    layer->getOutput(1)->setName("det_boxes");
    layer->getOutput(2)->setName("det_scores");
    layer->getOutput(3)->setName("det_classes");
    // 7
    network->markOutput(*layer->getOutput(0));
    network->markOutput(*layer->getOutput(1));
    network->markOutput(*layer->getOutput(2));
    network->markOutput(*layer->getOutput(3));

3.3、编译、测试代码

整合前面的代码，将yolov3.onnx 模型添加 shuffle、nms 层，导出 fp16 的 tensorrt 优化模型，代码如下

void build()
{
    std::string onnxFile = R"(E:\DeepLearning\darknet-yolo3-master\build\darknet\x64\yolov3.onnx)";

    Logger logger(nvinfer1::ILogger::Severity::kVERBOSE);

    /// 1.2 优化需要使用的临时变量
    // builder
    auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));

    // network
    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));

    // onnx parser
    auto parser = SampleUniquePtr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger));
    auto parsed = parser->parseFromFile(onnxFile.c_str(), (int)nvinfer1::ILogger::Severity::kVERBOSE);
     
    //  extend dims   
    //  scores [10647,80] -> [1, 10647,80],  boxex[10647,4] -> [1, 10647,4], 
    int nbOut = network->getNbOutputs(); // 2
    std::vector<nvinfer1::ITensor*> outputsTensors(2);
    outputsTensors[0] =  network->getOutput(1);
    outputsTensors[1] = network->getOutput(0);
    for(auto tensor : outputsTensors) {
        //auto dims = tensor->getDimensions();
        //auto name = tensor->getName();
        network->unmarkOutput(*tensor);
    } 
  
    std::vector<nvinfer1::ITensor*> newOutputsTensors = {
        unsqueezeTensor(network, outputsTensors[0], 0),
        unsqueezeTensor(network, outputsTensors[1], 0)
    };
    //for(auto tensor : newOutputsTensors) {
    //    auto dims = tensor->getDimensions();
    //    auto name = tensor->getName();
    //}

    // nms 
    nvinfer1:initLibNvInferPlugins(&logger, "");

    auto registry = nvinfer1::getBuilderPluginRegistry(nvinfer1::EngineCapability::kDEFAULT);
    auto creator = registry->getPluginCreator("EfficientNMS_TRT", "1");
        
    auto score_threshold = 0.25f;
    auto iou_threshold = 0.45f;
    auto max_output_boxes = 100;
    auto background_class = -1;
    auto score_activation = 0;
    auto class_agnostic = 1;
    auto box_coding = 1;
    auto plugin_version = "1";
    std::vector<nvinfer1::PluginField> new_pluginData_list = {
        {"score_threshold", &score_threshold, PluginFieldType::kFLOAT32, sizeof(score_threshold)},
        {"iou_threshold", &iou_threshold, PluginFieldType::kFLOAT32, sizeof(iou_threshold)},
        {"max_output_boxes", &max_output_boxes, PluginFieldType::kINT32, sizeof(max_output_boxes)},
        {"background_class", &background_class, PluginFieldType::kINT32, sizeof(background_class)},
        {"score_activation", &score_activation, PluginFieldType::kINT32, sizeof(score_activation)},
        {"class_agnostic", &class_agnostic, PluginFieldType::kINT32, sizeof(class_agnostic)},
        {"box_coding", &box_coding, PluginFieldType::kINT32, sizeof(box_coding)},
        //{"plugin_version", &plugin_version, PluginFieldType::kCHAR, sizeof(plugin_version)},
    };

	auto collection = nvinfer1::PluginFieldCollection();
    collection.nbFields = new_pluginData_list.size();
    collection.fields = new_pluginData_list.data();

    auto plugin = creator->createPlugin("nms_layer", &collection);

    auto layer = network->addPluginV2(newOutputsTensors.data(), newOutputsTensors.size(), *plugin);
    layer->getOutput(0)->setName("det_num");
    layer->getOutput(1)->setName("det_boxes");
    layer->getOutput(2)->setName("det_scores");
    layer->getOutput(3)->setName("det_classes");
    network->markOutput(*layer->getOutput(0));
    network->markOutput(*layer->getOutput(1));
    network->markOutput(*layer->getOutput(2));
    network->markOutput(*layer->getOutput(3));

    // config
    auto config = SampleUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    config->setFlag(nvinfer1::BuilderFlag::kFP16);

    auto profileStream = makeCudaStream();
    config->setProfileStream(*profileStream);

    /// 1.3 编译优化engine并序列化
// serialize
    SampleUniquePtr<nvinfer1::IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};

    std::ofstream ofs(R"(E:\DeepLearning\darknet-yolo3-master\build\darknet\x64\yolov3.onnx.nms.trt)", std::ostream::binary);
    ofs.write(static_cast<const char*>(plan->data()), plan->size());
    ofs.close();
}

之后使用上一节 2.4、Tensorrt API c++测试 yolo_nms 中的测试代码进行测试，一切正常。

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.coloradmin.cn/o/1876279.html

如若内容造成侵权/违法违规/事实不符，请联系多彩编程网进行投诉反馈，一经查实，立即删除！

相关文章

墨刀原型--多tab切换显示对应页面场景交互步骤

墨刀原型--多tab切换显示对应页面场景交互步骤

一般我们画原型页面&#xff0c;PC端或者APP端或小程序端&#xff0c;都会有页面会切换多个tab或状态&#xff0c;同时对应页面显示对应的页面数据。 设计思路如下&#xff1a; 以订单列表页面为例&#xff1a; 可以将订单列表页面分为3部分&#xff0c;固定的头部、状态栏、…
阅读更多...
驾校预约小程序系统的设计

驾校预约小程序系统的设计

管理员账户功能包括&#xff1a;系统首页&#xff0c;个人中心&#xff0c;学员管理&#xff0c;教练管理&#xff0c;驾校信息管理&#xff0c;驾校车辆管理&#xff0c;教练预约管理&#xff0c;考试信息管理 微信端账号功能包括&#xff1a;系统首页&#xff0c;驾校信息&am…
阅读更多...
Python operator模块这么用，效率杠杠的！

Python operator模块这么用，效率杠杠的！

目录 1、基础操作符应用 🐍 1.1 加载operator模块 1.2 使用itemgetter进行排序 1.3 attrgetter与方法调用 2、高级功能探索 🔍 2.1 methodcaller的妙用 2.2 操作符重载与定制 3、结合lambda表达式 ✨ 3.1 lambda与operator模块协同工作 3.2 实战案例分析 4、结合…
阅读更多...
面试-javaIO机制

面试-javaIO机制

1.BIO BIO&#xff1a;是传统的javaIO以及部分java.net下部分接口和类。例如&#xff0c;socket,http等&#xff0c;因为网络通信同样是IO行为。传统IO基于字节流和字符流进行操作。提供了我们最熟悉的IO功能&#xff0c;譬如基于字节流的InputStream 和OutputStream.基于字符流…
阅读更多...
DataWhale-吃瓜教程学习笔记 (五）

DataWhale-吃瓜教程学习笔记 (五）

学习视频&#xff1a;第4章-决策树_哔哩哔哩_bilibili 西瓜书对应章节&#xff1a; 第四章 4.1&#xff1b;4.2 决策树算法原理 - 逻辑角度 if...else.. 语句的组合&#xff0c;不断的选择 - 几何角度 根据某种准则划分特征空间 最终目的&#xff1a;提高分类样本的纯度 I…
阅读更多...
Qt开发笔记：Qt3D三维开发笔记（一）：Qt3D三维开发基础概念介绍

Qt开发笔记：Qt3D三维开发笔记（一）：Qt3D三维开发基础概念介绍

若该文为原创文章&#xff0c;转载请注明原文出处 本文章博客地址&#xff1a;https://blog.csdn.net/qq21497936/article/details/140059315 长沙红胖子Qt&#xff08;长沙创微智科&#xff09;博文大全&#xff1a;开发技术集合&#xff08;包含Qt实用技术、树莓派、三维、O…
阅读更多...
java第二十九课 —— 断点 | 零钱通项目

java第二十九课 —— 断点 | 零钱通项目

断点调试&#xff08;debug&#xff09; 实际需求 在开发中&#xff0c;新手程序员在查找错误时&#xff0c;这时老程序员就会温馨提示&#xff0c;可以用断点调试步一步的看源码执行的过程&#xff0c;从而发现错误所在。 重要提示&#xff1a;在断点调试过程中&#xff0c;…
阅读更多...
如何安装和卸载软件？

如何安装和卸载软件？

如何安装和卸载软件&#xff1f; &#x1f4bb; 如何安装和卸载软件&#xff1f;——默语的详细教程摘要引言正文内容&#x1f5a5;️ 在Windows上安装和卸载软件安装软件卸载软件 &#x1f34f; 在Mac上安装和卸载软件安装软件卸载软件 &#x1f914; QA环节&#x1f4dd; 表格…
阅读更多...
【nvm】如何使用nvm优雅的管理Node.js

【nvm】如何使用nvm优雅的管理Node.js

希望文章能给到你启发和灵感&#xff5e; 如果觉得文章对你有帮助的话&#xff0c;点赞 关注 收藏 支持一下博主吧&#xff5e; 阅读指南 开篇说明一、基础环境说明1.1 硬件环境1.2 软件环境 二、什么是nvm?2.1 概念2.1 安装2.1.1 对于Mac系统2.1.2 对于Windows系统2.1.3 对于…
阅读更多...
完全离线的本地问答模型LocalGPT如何实现无公网IP远程连接提问

完全离线的本地问答模型LocalGPT如何实现无公网IP远程连接提问

文章目录 前言环境准备1. localGPT部署2. 启动和使用3. 安装cpolar 内网穿透4. 创建公网地址5. 公网地址访问6. 固定公网地址 前言 本文主要介绍如何本地部署LocalGPT并实现远程访问&#xff0c;由于localGPT只能通过本地局域网IP地址端口号的形式访问&#xff0c;实现远程访问…
阅读更多...
基于OrangePi AIpro + owncloud 5分钟搭建一个私有网盘

基于OrangePi AIpro + owncloud 5分钟搭建一个私有网盘

OrangePi AIpro自带镜像系统已预装了docker&#xff0c;这里我们直接基于docker安装owncloud。 准备 切换用户&#xff1a; HwHiAiUser 默认密码&#xff1a;Mind123 su HwHiAiUser 创建文件夹 sudo mkdir /home/SummerGao/owncloud-docker-server 切换至刚创建的文件夹下…
阅读更多...
使用 Compose Multiplatform Media Player 实现跨平台媒体播放

使用 Compose Multiplatform Media Player 实现跨平台媒体播放

使用 Compose Multiplatform Media Player 实现跨平台媒体播放 在跨平台开发中，媒体播放功能是一个常见且重要的需求。Compose Multiplatform Media Player 是一个专为 Compose Multiplatform 项目设计的强大媒体播放器库，它可以在 iOS 和 Android 平台上无缝实现视频播放、…
阅读更多...
Prism 目录方式加载插件，提示`xxx.resources.dll`找不到

Prism 目录方式加载插件，提示`xxx.resources.dll`找不到

问题场景 前置条件 使用Prism 实现的目录配置方式加载插件&#xff1b; 有两个模块插件&#xff1a;ModuleA 以及 MouduleB。 问题现象 独立项目作为插件被加载时&#xff0c;加载指定模块中的用户控件&#xff0c;程序能正常运行&#xff0c;点击模块进行加载。 但输出窗…
阅读更多...
cartographer从入门到精通（一）：cartographer介绍

cartographer从入门到精通（一）：cartographer介绍

一、cartographer重要文档 有关cartographer的资料有2个比较重要的网站&#xff0c;我们的介绍也是基于这两个网站&#xff0c;其中会加入自己的一些理解&#xff0c;后续也有一些对代码的修改&#xff0c;来实现我们想完善的功能。 1-Cartographer 2-Cartographer ROS 第1个…
阅读更多...
Python 面试【★★★★★】

Python 面试【★★★★★】

欢迎莅临我的博客 &#x1f49d;&#x1f49d;&#x1f49d;&#xff0c;很高兴能够在这里和您见面&#xff01;希望您在这里可以感受到一份轻松愉快的氛围&#xff0c;不仅可以获得有趣的内容和知识&#xff0c;也可以畅所欲言、分享您的想法和见解。 推荐:「stormsha的主页」…
阅读更多...
STM32第九课：DHT11温湿度传感器

STM32第九课：DHT11温湿度传感器

文章目录 需求一、DHT11温湿度传感器二、模块配置流程1.配置时钟和IO2.读取数据3.数据处理 三、导入语音模块四、关键代码总结 需求 1.完成DHT11温湿度检测模块的配置。 2.处理DHT11获取的数据&#xff0c;在串口打印处理后的实时数据。 2.通过Su-03t语音识别模块实现实时温湿…
阅读更多...
WLAN 4-Way Handshake如何生成GTK？

WLAN 4-Way Handshake如何生成GTK？

关于Wi-Fi的加密认证过程&#xff0c;可以参考如下链接&#xff0c;今天我们来理解如何生成GTK。 WLAN数据加密机制_tls加密wifi-CSDN博客 1 GTK GTK&#xff08;Group Temporal Key&#xff09;是由AP通过GMK生成&#xff0c;长度为128位&#xff0c;并在四次握手的第三步中…
阅读更多...
哈尔滨高校大学智能制造实验室数字孪生可视化系统平台项目的验收

哈尔滨高校大学智能制造实验室数字孪生可视化系统平台项目的验收

哈尔滨高校大学智能制造实验室数字孪生可视化系统平台项目的验收&#xff0c;标志着这一技术在教育领域的应用取得了新的突破。项目旨在开发一个数字孪生可视化系统平台&#xff0c;用于哈尔滨高校大学智能制造实验室的设备模拟、监测与数据分析。项目的主要目标包括&#xff1…
阅读更多...
数据结构_线性表

数据结构_线性表

线性表的定义和特点 线性表是具有相同特性的数据元素的一个有限序列 :线性起点/起始节点 :的直接前驱 :的直接后继 :线性终点/终端节点 n:元素总个数,表长 下标:是元素的序号,表示元素在表中的位置 n0时称为空表 线性表 由n(n>0)个数据元素(结点),组成的有限序列 将…
阅读更多...
Quantlab5.0：一切围绕可实盘策略驱动开发

Quantlab5.0：一切围绕可实盘策略驱动开发

原创文章第573篇&#xff0c;专注“AI量化投资、世界运行的规律、个人成长与财富自由"。 2024年上半年即将结束&#xff0c;开始准备星球下半年的工作。 目前设想的——Quantlab5.0&#xff0c;之所以升级一个大版本&#xff0c;与4.x有很大不同。 5.0专注策略开发&…
阅读更多...
推荐文章
最新文章

Copyright @ 2022~2023