上饶网站开发,室内设计师第一网站,宁波网站设计首选荣盛网络,给宝宝做辅食的网站之前搞过不少部署#xff0c;也玩过tensorRT部署模型#xff08;但都是模型推理用gpu#xff0c;后处理还是用cpu进行#xff09;#xff0c;有网友问能出一篇tensorRT用gpu对模型后处理进行加速的。由于之前用的都是非cuda支持的边缘芯片#xff0c;没有写过cuda代码也玩过tensorRT部署模型但都是模型推理用gpu后处理还是用cpu进行有网友问能出一篇tensorRT用gpu对模型后处理进行加速的。由于之前用的都是非cuda支持的边缘芯片没有写过cuda代码这个使得我犹豫不决。零零总总恶补了一点cuda编程于是就有了这篇博客【yolov11 部署 TensorRT预处理和后处理用 C cuda 加速速度快到飞起】。既然用cuda实现就不按照部署其他芯片比如rk3588那套导出onnx的流程修改导出onnx代码要尽可能多的操作都在cuda上运行这样导出onnx的方式就比较简单yolov11官方导出onnx的方式。 rtx4090显卡、模型yolov11n输入分辨率640x64080个类别、量化成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,480,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,8400AnchorCount: 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(modelyolov11n.pt) # load a pretrained model (recommended for training)
results model(taskdetect, sourcer./bus.jpg, saveTrue) # predict on an imagemodel.export(formatonnx, imgsz640, simplifyTrue)
编译 修改 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输入分辨率640x64080个类别、量化成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::durationfloat, 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;
预处理用cuda加速 代码中已实现用CUDA_npp_LIBRARY进行预处理如果有环境可以打开进一步加速修改位置CMakelist.txt 已进行了注释、用CPU或GPU预处理打开对应的宏 #define USE_GPU_PREPROCESS 1) 重新搭建了一个支持用gpu对预处理进行加速的环境rtx4090显卡、模型yolov11n输入分辨率640x64080个类别、量化成FP16模型。对比结果如下这台机器相比上面贴图中时耗更短可能是这台机器的cpu性能比较强。以下给出的时耗是预处理模型推理后处理。
cpu做预处理模型推理cpu做后处理 cpu做预处理模型推理gpu做后处理 gpu做预处理gpu做后处理
后续优化点
1、把nms过程也用cuda实现参加nms的框不多但也是一个优化点持续更新中
