??之前搞過不少部署，也玩過tensorRT部署模型（但都是模型推理用gpu，后處理還是用cpu進行），有網(wǎng)友問能出一篇tensorRT用gpu對模型后處理進行加速的。由于之前用的都是非cuda支持的邊緣芯片，沒有寫過cuda代碼，這個使得我猶豫不決。零零總總惡補了一點cuda編程，于是就有了這篇博客【yolov11 部署 TensorRT，預(yù)處理和后處理用 C++ cuda 加速，速度快到飛起】。既然用cuda實現(xiàn)，就不按照部署其他芯片（比如rk3588）那套導出onnx的流程修改導出onnx代碼，要盡可能多的操作都在cuda上運行，這樣導出onnx的方式就比較簡單（yolov11官方導出onnx的方式）。

??rtx4090顯卡、模型yolov11n（輸入分辨率640x640，80個類別）、量化成FP16模型，最快1000fps

??本示例中，包含完整的代碼、模型、測試圖片、測試結(jié)果。

??后處理部分用cuda 核函數(shù)實現(xiàn)，并不是全部后處理都用cuda實現(xiàn)【cuda實現(xiàn)后處理代碼】；純cpu實現(xiàn)后處理部分代碼分支【cpu實現(xiàn)后處理代碼】

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

cuda 核函數(shù)主要做的操作

??由于按照yolov11官方導出的onnx，后處理需要做的操作只有nms了。mns流程：選出大于閾值的框，在對這些大于閾值的框進行排序，在進行nms操作。這幾部中最耗時的操作（對類別得分選出最大對應(yīng)的類別，判斷是否大于閾值）是選出所有大于閾值的框，經(jīng)過這一步后實際參加nms的框沒有幾個（比如圖像中有30個目標，每個目標出來了15個框，也才450個框），因此主要對這一步操作“選出所有大于閾值的框”用cuda實現(xiàn)。當然后續(xù)還可以繼續(xù)優(yōu)化，把nms的過程用cuda核函數(shù)進行實現(xiàn)。

??核函數(shù)的實現(xiàn)如下：模型輸出維度（1,（4+80）,8400），主要思想流程用8400個線程，實現(xiàn)對80個類別選出最大值，并判斷是否大于閾值。

__global__ void GetNmsBeforeBoxesKernel(float *SrcInput, int AnchorCount, int ClassNum, float ObjectThresh, int NmsBeforeMaxNum, DetectRect* OutputRects, int *OutputCount)
{/***功能說明：用8400個線程，實現(xiàn)對80個類別選出最大值，并判斷是否大于閾值，把大于閾值的框記錄下來后面用于參加mnsSrcInput: 模型輸出（1,84,8400）AnchorCount: 8400ClassNum: 80ObjectThresh: 目標閾值（大于該閾值的目標才輸出）NmsBeforeMaxNum: 輸入nms檢測框的最大數(shù)量，前面申請的了一塊兒顯存來裝要參加nms的框，防止越界OutputRects: 大于閾值的目標框OutputCount: 大于閾值的目標框個數(shù)***/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 對應(yīng)的TensorRT位置

cd yolov11_tensorRT_postprocess_cuda
mkdir build
cd build
cmake ..
make

運行

# 運行時如果.trt模型存在則直接加載，若不存會自動先將onnx轉(zhuǎn)換成 trt 模型，并存在給定的位置，然后運行推理。
cd build
./yolo_trt

測試效果

onnx 測試效果

tensorRT 測試效果

tensorRT 時耗（cuda實現(xiàn)部分后處理）

??示例中用cpu對圖像進行預(yù)處理（由于本臺機器搭建的環(huán)境不匹配，不能用cuda對預(yù)處理進行加速）、用rtx4090顯卡進行模型推理、用cuda對后處理進行加速。使用的模型yolov11n（輸入分辨率640x640，80個類別）、量化成FP16模型。以下給出的時耗是：預(yù)處理+模型推理+后處理。

cpu做預(yù)處理+模型推理+gpu做后處理

tensorRT 時耗（純cpu實現(xiàn)后處理）【cpu實現(xiàn)后處理代碼分支】

cpu做預(yù)處理+模型推理+cpu做后處理

替換模型說明

??修改相關(guān)的路徑

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;

預(yù)處理用cuda加速

??代碼中已實現(xiàn)用CUDA_npp_LIBRARY進行預(yù)處理，如果有環(huán)境可以打開進一步加速（修改位置：CMakelist.txt 已進行了注釋、用CPU或GPU預(yù)處理打開對應(yīng)的宏 #define USE_GPU_PREPROCESS 1）)

??重新搭建了一個支持用gpu對預(yù)處理進行加速的環(huán)境：rtx4090顯卡、模型yolov11n（輸入分辨率640x640，80個類別）、量化成FP16模型。對比結(jié)果如下：這臺機器相比上面貼圖中時耗更短，可能是這臺機器的cpu性能比較強。以下給出的時耗是：預(yù)處理+模型推理+后處理。

cpu做預(yù)處理+模型推理+cpu做后處理

cpu做預(yù)處理+模型推理+gpu做后處理

gpu做預(yù)處理+gpu做后處理

后續(xù)優(yōu)化點

1、把nms過程也用cuda實現(xiàn)，參加nms的框不多，但也是一個優(yōu)化點，持續(xù)更新中