当前位置：首页 > news >正文

yolov11 部署 TensorRT，预处理和后处理用 C++ cuda 加速，速度快到飞起

news 2025/7/14 13:52:31

之前搞过不少部署，也玩过tensorRT部署模型（但都是模型推理用gpu，后处理还是用cpu进行），有网友问能出一篇tensorRT用gpu对模型后处理进行加速的。由于之前用的都是非cuda支持的边缘芯片，没有写过cuda代码，这个使得我犹豫不决。零零总总恶补了一点cuda编程，于是就有了这篇博客【yolov11 部署 TensorRT，预处理和后处理用 C++ cuda 加速，速度快到飞起】。既然用cuda实现，就不按照部署其他芯片（比如rk3588）那套导出onnx的流程修改导出onnx代码，要尽可能多的操作都在cuda上运行，这样导出onnx的方式就比较简单（yolov11官方导出onnx的方式）。

rtx4090显卡、模型yolov11n（输入分辨率640x640，80个类别）、量化成FP16模型，最快1000fps

本示例中，包含完整的代码、模型、测试图片、测试结果。

后处理部分用cuda 核函数实现，并不是全部后处理都用cuda实现【cuda实现后处理代码】；纯cpu实现后处理部分代码分支【cpu实现后处理代码】

使用的TensorRT版本：TensorRT-8.6.1.6
cuda:11.4
显卡：RTX4090

cuda 核函数主要做的操作

由于按照yolov11官方导出的onnx，后处理需要做的操作只有nms了。mns流程：选出大于阈值的框，在对这些大于阈值的框进行排序，在进行nms操作。这几部中最耗时的操作（对类别得分选出最大对应的类别，判断是否大于阈值）是选出所有大于阈值的框，经过这一步后实际参加nms的框没有几个（比如图像中有30个目标，每个目标出来了15个框，也才450个框），因此主要对这一步操作“选出所有大于阈值的框”用cuda实现。当然后续还可以继续优化，把nms的过程用cuda核函数进行实现。

核函数的实现如下：模型输出维度（1,（4+80）,8400），主要思想流程用8400个线程，实现对80个类别选出最大值，并判断是否大于阈值。

__global__ void GetNmsBeforeBoxesKernel(float *SrcInput, int AnchorCount, int ClassNum, float ObjectThresh, int NmsBeforeMaxNum, DetectRect* OutputRects, int *OutputCount)
{/***功能说明：用8400个线程，实现对80个类别选出最大值，并判断是否大于阈值，把大于阈值的框记录下来后面用于参加mnsSrcInput: 模型输出（1,84,8400）AnchorCount: 8400ClassNum: 80ObjectThresh: 目标阈值（大于该阈值的目标才输出）NmsBeforeMaxNum: 输入nms检测框的最大数量，前面申请的了一块儿显存来装要参加nms的框，防止越界OutputRects: 大于阈值的目标框OutputCount: 大于阈值的目标框个数***/int ThreadId = blockIdx.x * blockDim.x + threadIdx.x;if (ThreadId >= AnchorCount){return;}float* XywhConf = SrcInput + ThreadId;float CenterX = 0, CenterY = 0, CenterW = 0, CenterH = 0;float MaxScore = 0;int MaxIndex = 0;DetectRect TempRect;for (int j = 4; j < ClassNum + 4; j ++) {if (4 == j){MaxScore = XywhConf[j * AnchorCount];MaxIndex = j;   } else {if (MaxScore <  XywhConf[j * AnchorCount]){MaxScore = XywhConf[j * AnchorCount];MaxIndex = j;   }}  }if (MaxScore > ObjectThresh){int index = atomicAdd(OutputCount, 1);if (index > NmsBeforeMaxNum){return;}CenterX = XywhConf[0 * AnchorCount];CenterY = XywhConf[1 * AnchorCount];CenterW = XywhConf[2 * AnchorCount];CenterH = XywhConf[3 * AnchorCount ];TempRect.classId = MaxIndex - 4;TempRect.score = MaxScore;TempRect.xmin = CenterX - 0.5 * CenterW;TempRect.ymin = CenterY - 0.5 * CenterH;TempRect.xmax = CenterX + 0.5 * CenterW;TempRect.ymax = CenterY + 0.5 * CenterH;OutputRects[index] = TempRect;}
}

导出onnx模型

按照yolov11官方导出的方式如下：

from ultralytics import YOLO
model = YOLO(model='yolov11n.pt')  # load a pretrained model (recommended for training)
results = model(task='detect', source=r'./bus.jpg', save=True)  # predict on an imagemodel.export(format="onnx", imgsz=640, simplify=True)

编译

修改 CMakeLists.txt 对应的TensorRT位置
在这里插入图片描述

cd yolov11_tensorRT_postprocess_cuda
mkdir build
cd build
cmake ..
make

运行

# 运行时如果.trt模型存在则直接加载，若不存会自动先将onnx转换成 trt 模型，并存在给定的位置，然后运行推理。
cd build
./yolo_trt

测试效果

onnx 测试效果

在这里插入图片描述

tensorRT 测试效果

在这里插入图片描述

tensorRT 时耗（cuda实现部分后处理）

示例中用cpu对图像进行预处理（由于本台机器搭建的环境不匹配，不能用cuda对预处理进行加速）、用rtx4090显卡进行模型推理、用cuda对后处理进行加速。使用的模型yolov11n（输入分辨率640x640，80个类别）、量化成FP16模型。以下给出的时耗是：预处理+模型推理+后处理。

cpu做预处理+模型推理+gpu做后处理
在这里插入图片描述

tensorRT 时耗（纯cpu实现后处理）【cpu实现后处理代码分支】

cpu做预处理+模型推理+cpu做后处理

替换模型说明

修改相关的路径

std::string OnnxFile = "/root/autodl-tmp/yolov11_tensorRT_postprocess_cuda/models/yolov11n.onnx";std::string SaveTrtFilePath = "/root/autodl-tmp/yolov11_tensorRT_postprocess_cuda/models/yolov11n.trt";cv::Mat SrcImage = cv::imread("/root/autodl-tmp/yolov11_tensorRT_postprocess_cuda/images/test.jpg");int img_width = SrcImage.cols;int img_height = SrcImage.rows;std::cout << "img_width: " << img_width << " img_height: " << img_height << std::endl;CNN YOLO(OnnxFile, SaveTrtFilePath, 1, 3, 640, 640);auto t_start = std::chrono::high_resolution_clock::now();int Temp = 2000;int SleepTimes = 0;for (int i = 0; i < Temp; i++){YOLO.Inference(SrcImage);std::this_thread::sleep_for(std::chrono::milliseconds(SleepTimes));}auto t_end = std::chrono::high_resolution_clock::now();float total_inf = std::chrono::duration<float, std::milli>(t_end - t_start).count();std::cout << "Info: " << Temp << " times infer and gpu postprocess ave cost: " << total_inf / float(Temp) - SleepTimes << " ms." << std::endl;