TensorRT 使用 - 以 yolov9 為例
===
# 摘要
* 本文紀錄了 tensorRT 的使用並以 yolov9 當作範例來去實踐, 在 v10 已經有 NMS free 達成但NMS 這一經典的方法覺得還是該寫進此筆記作為紀錄, 有需要也可以回頭複習
* 後續也可基於此專案改寫不同的推論引擎
# 環境建制
## 硬體環境 ：
* Dell G15
* GPU : RTX3050 Laptop 4G
* CPU : 12th Gen Intel® Core™ i5-12500H, 16 cores
* Memory : 16GB
## 軟體套件 :
* opencv : 4.10
* tensorRT : release/10.6
* cuda : 12.6
* yolov9
## 編譯方式 ：
* 目前相關程式紀錄在 [github](https://github.com/GU-Lin/AI_engine.git), 首先使用 make_build.sh 這個 script 檔指令如下, 會使用 docker 來建立相關的 library 預設會是建立在 build_x86_64 的資料夾下
* 執行完後可在 builx_x86_64 底下看到 opencv、tensorRT的資料夾
```script=
./make_build.sh --opencv --tensorRT
```
# TensorRT 使用
## 模型轉換
* TensorRT 使用的檔案為 .trt, 而這部份可透過將 .onnx 轉換成 .trt 達成過程中也可根據想要答標的速度或是精度做選擇（例如 fp32、fp16、int8等）,
* 然而 .trt 會根據當前硬體條件（也就是 GPU的算例）做對應的轉換, 以自身筆電為例所使用的算力為 Compute Capability : 8.6 如果放到 7.x 或其他版本的硬體上此 .trt 則無法是用
* Compute Capability 可在 https://developer.nvidia.com/cuda-gpus 下查詢, 或是轉換時即可在資訊中顯示
* 執行以下指令即可將 .onnx 轉為 .trt, yolov9-c-converted.onnx可到 yolov9 的 github 下載模型並自行轉成 onnx
```bash=
cd build_x86_64
export LD_LIBRARY_PATH=`pwd`/lib
./bin/trtexec --onnx=../../data/yolov9-c-converted.onnx --saveEngine=../../yolov9-c-converted.trt
```
## 模型讀取
* 在讀取時有以下物件需要使用, 讀取模型的方式放在 trtInferenceTool.cpp 內
* IExecutionContext : 用於執行實際推論
* ICudaEngine : 提供創建上下文的方法及表示靜態的模型推理
* IRuntime : 反序列化模型
```cpp=
std::unique_ptr<IExecutionContext> m_context;
std::shared_ptr<ICudaEngine> m_engine;
std::unique_ptr<IRuntime> m_runtime;
```
* 首先給定路徑讀取模型
```cpp=
std::ifstream file(path, std::ios::binary);
if (!file.is_open()) {
std::cerr << "Error opening engine file: " << path << std::endl;
return ;
}
// 讀取文件內容
file.seekg(0, file.end);
size_t size = file.tellg();
file.seekg(0, file.beg);
std::vector<char> engineData(size);
file.read(engineData.data(), size);
file.close();
```
* 創建 runtime
```cpp=
// 創建 TensorRT runtime
m_runtime.reset(createInferRuntime(gLogger));
if (!m_runtime) {
std::cerr << "Failed to create InferRuntime" << std::endl;
return ;
}
```
* 創建 engine 即將模型反序列化, 反序列化是只將模型轉換為物件的過程
* 舉例 : 讀取 .json 檔將檔案內容反序列化成字典的形式做使用
```cpp=
m_engine = std::shared_ptr<ICudaEngine>(
m_runtime->deserializeCudaEngine(engineData.data(), size), samplesCommon::InferDeleter());
if (!m_engine) {
std::cerr << "Failed to deserialize CUDA engine" << std::endl;
return ;
}
```
* 創建 context, 可透過 getIOTensorName 來得到 input/output 的名稱, 可用於獲取模型輸入輸出的資訊
```cpp=
m_context.reset(m_engine->createExecutionContext());
inputName = m_engine->getIOTensorName(0);
outputName = m_engine->getIOTensorName(1);
std::cout << "Input Name : " << inputName << std::endl;
std::cout << "Output Name : " << outputName << std::endl;
```
* 記憶體分配 : 在此透過 cuda 分配記憶體給 GPU 目的是讓 .trt 使用特定 GPU 記憶體位置進行推論加速
* buffers[0] 為 input, buffers[1] 為 output
```cpp=
m_inputSize = getIOSize(inputName);
m_outputSize = getIOSize(outputName);
cudaMalloc(&buffers[0],m_inputSize*sizeof(float));
cudaMalloc(&buffers[1],m_outputSize*sizeof(float));
m_context->setTensorAddress(inputName,buffers[0]);
m_context->setTensorAddress(outputName,buffers[1]);
hostData = new float[m_outputSize];
std::cout << "Input size is " << m_inputSize << std::endl;
std::cout << "Output size is " << m_outputSize << std::endl;
std::cout << "Load " << path << " successful" << std::endl;
```
# yolov9 模型的 IO
* Input : 為一個 [1,3,640,640] 的彩色三通道圖片, 但需要經過 normalize 到 [0,1] 之間且轉為 RGB 應向 CHW 的格式
* Output : 為 [1,84,8400] 的樣子, 而其中輸出數字代表的意義如下
* Cx : Bounding Box 的中心象素 x 座標
* Cy : Bounding Box 的中心象素 y 座標
* W : Bounding Box 的寬
* H : Bounding Box 的高
* 以第一直行(column)的方式來看, 第一個物件前4項就是 Bounding box 的位置, 後面 80 項就是對應類別的分數, 如果 Class 10 的分數最高則第一個物件就是被判定為類別 10, 而從輸出可以得到 yolov9 可以一口氣偵測 8400 個物件, 附上 [yolov9 類別對應表](https://github.com/WongKinYiu/yolov9/blob/main/data/coco.yaml)
# yolov9 c++ 實現
## Preprocess
* 預處理的過程為了讓每張影像都縮放成 640 x 640 並使用 warpAffine 的方式進行縮放, 而縮放的過程中如果還有不足的邊框(例如縮放完變成 640 x 512) 則圖片會置中並把多餘的剩下的象素補邊, 象素內容為 (114,114,114)
```cpp=
void objectDetection::preprocess(cv::Mat input, cv::Mat &output)
{
cv::Mat src = input.clone();
cv::Size sizeDst(640,640);
cv::Size sizeSrc = src.size();
float ratioWidth = float(sizeDst.width)/float(sizeSrc.width);
float ratioHeight = float(sizeDst.height)/float(sizeSrc.height);
float r = ratioWidth < ratioHeight ? ratioWidth : ratioHeight;
cv::Size sizeEnlarge(std::round(r*sizeSrc.width),std::round(r*sizeSrc.height));
// Compute padding
int dw = (sizeDst.width - sizeEnlarge.width)/2;
int dh = (sizeDst.height - sizeEnlarge.height)/2;
cv::Mat M = cv::Mat::zeros(2,3,CV_32F);
M.at<float>(0,0) = r;
M.at<float>(1,1) = r;
M.at<float>(0,2) = dw;
M.at<float>(1,2) = dh;
cv::warpAffine(src,src,M,(cv::Size(640,640)),cv::INTER_LINEAR,cv::BORDER_CONSTANT,cv::Scalar(114,114,114));
src.copyTo(output) ;
}
```
* 再來需要將影像轉成正確的通道
```cpp=
void objectDetection::HWC2NormalCHW(cv::Mat input, std::vector<float> &data)
{
std::vector<cv::Mat> inputChannel(3);
cv::Mat img;
cv::cvtColor(input, img, cv::COLOR_BGR2RGB);
img.convertTo(img, CV_32F, 1.0f/255.f);
cv::split(img, inputChannel);
std::vector<float>result;
for(int i = 0; i < 3; i++)
{
std::vector<float> temp = std::vector<float>(inputChannel[i].reshape(1, 1));
result.insert(result.end(),temp.begin(),temp.end());
}
data = result;
}
```
## Inference
* 在做推論時最常需要做的就是 cpu 與 gpu 之間資料的拷貝, 具體流程如下
* Host to Device : cpu -> gpu
* AI inference : m_context 推論
* Device to Host : gpu -> cpu
```cpp=
void TRTInferenceTool::run(std::vector<float> &input, cv::Mat &output)
{
cudaStream_t stream;
cudaStreamCreate(&stream);
// From host to device
cudaMemcpyAsync(buffers[0], input.data(), m_inputSize * sizeof(float), cudaMemcpyHostToDevice, stream);
// Enqueue excute
m_context->enqueueV3(stream);
// From device to host
cudaMemcpyAsync(hostData, buffers[1], m_outputSize * sizeof(float), cudaMemcpyDeviceToHost, stream);
cudaStreamSynchronize(stream);
cv::Mat m(m_outputClass,m_outputBoxNum, CV_32FC1, hostData);
output = m.clone();
}
```
## Postprocess
* 創建了以下 struct 紀錄 bounding box
```cpp=
struct PredBox
{
float cx;
float cy;
float width;
float height;
float score;
int label;
};
```
* 後處理分成以下三個部分
* 挑選是否過 class 閥值的物件框
* 物件框進行排序
* NMS 將重複框消除
## 挑選是否過 class 閥值的物件框
* 可看到有使用 minMaxLoc 來找尋最大的類別及對應的類別分數, 如果分數有過閥值則紀錄下來到 m_PredBox_vector 內
```cpp=
for(int i = 0; i < input.cols; i++)
{
cv::Mat res = input.col(i);
float cx = res.at<float>(0);
float cy = res.at<float>(1);
float w = res.at<float>(2);
float h = res.at<float>(3);
cv::Mat cls = res(cv::Range(4,84),cv::Range::all());
cv::Point minIdx, maxIdx;
double minVal, maxVal;
cv::minMaxLoc(cls, &minVal, &maxVal, &minIdx, &maxIdx);
if(maxVal > m_conf_thres)
{
PredBox pbox;
pbox.cx = cx;
pbox.cy = cy;
pbox.width = w;
pbox.height = h;
pbox.score = maxVal;
pbox.label = maxIdx.y;
m_PredBox_vector.push_back(pbox);
}
// break;
}
```
## 物件框進行排序
* 根據比分由大至小進行排序
```cpp=
if(m_PredBox_vector.size()>0)
{
std::sort(m_PredBox_vector.begin(), m_PredBox_vector.end(), [](PredBox a, PredBox b)
{
return a.score > b.score;
});
}
```
## NMS 將重複框消除
* NMS 會將兩個物件框做 IOU 的計算, 如果 IOU 重疊比例過閥值則會跳過不紀錄, 紀錄的方式是用一 ```vector<bool> table ```的方式來看是否需要跳過此框
```cpp=
void objectDetection::NMS(std::vector<PredBox> &pred)
{
std::vector<PredBox> result;
std::vector<bool> table(pred.size(),false);
for(int i = 0; i < pred.size(); i++)
{
if(table[i])continue;
PredBox temp = pred[i];
result.push_back(temp);
for(int j = i+1; j < pred.size(); j++)
{
if(table[j])continue;
if(iou(temp,pred[j]) > m_iou_thres)
{
table[j] = true;
}
}
}
pred = result;
}
```
# 畫框
* 目的需要將物件框畫到原圖上而不是補上邊的圖, 所以會計算縮放的比例大小即位移多寡
* 計算完後用 cv::rectangle 畫取 bounding box
* 寫上類別後利用 cv::getTextSize 來得到畫上的字串佔據多少的寬高, 目的是自動讓填滿可視化的所呈現的字體周邊象素, 不然有的字長是 10 有的是 3 不太好調整一個通用值
```cpp=
void objectDetection::draw(cv::Mat &input, cv::Mat &output, std::vector<PredBox> &pred)
{
output = input.clone();
float ration_w = 640/float(input.cols);
float ration_h = 640/float(input.rows);
for(int i = pred.size()-1; i >= 0 ; i--)
{
int left = int(pred[i].cx-pred[i].width/2);
int top = int(pred[i].cy-pred[i].height/2);
int width = pred[i].width;
int height = pred[i].height;
if(ration_h > ration_w)
{
left = left /ration_w;
top = (top-(640-ration_w*input.rows)/2)/ration_w;
width = width / ration_w;
height = height / ration_w;
}else
{
left = left /ration_h;
top = (top-(640-ration_h*input.cols)/2)/ration_h;
width = width / ration_h;
height = height / ration_h;
}
// Rectangle
cv::rectangle(output, cv::Point(left,top), cv::Point(left + width,top + height), cv::Scalar(0,255,0),2);
// Name
std::string class_string = coconame[pred[i].label] + ' ' + std::to_string(pred[i].score).substr(0, 4);
int baseline = 0;
int fontFace = cv::FONT_HERSHEY_SIMPLEX; // 字體類型
double fontScale = 1.0; // 字體大小
int thickness = 2; // 字體粗細
cv::Size textSize = cv::getTextSize(class_string, fontFace, fontScale, thickness, &baseline);
cv::Rect textRect(left, top- textSize.height, textSize.width, textSize.height);
cv::rectangle(output,textRect,cv::Scalar(0,255,0),cv::FILLED);
putText(output, class_string, cv::Point(left , top), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 1, 0);
}
}
```
# 編譯程式並執行
```cpp=
cmake CMakeList.txt
make
./main
```
# 實際輸出
source :
https://github.com/WongKinYiu/yolov9
https://blog.csdn.net/qq_40672115/article/details/136492338
github : https://github.com/GU-Lin/AI_engine
###### tags : `AI`, `AI 部屬`
Sign in with Wallet
Connect another wallet