做網(wǎng)站人網(wǎng)頁設(shè)計制作網(wǎng)站素材
??之前搞過不少部署,也玩過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ù)更新中