在瞭解一些概念之前一直看不懂上交22年開源的TRTModule.cpp和.hpp，好在交爺寫的足夠模組化，可以配好環境開箱即用，移植很簡單。最近稍微瞭解了神經網路的一些概念，又看了TensorRT的一些api，遂試著部署一下自己在MNIST手寫數字資料集上訓練的一個LeNet模型，識別率大概有98.9%，實現用pytorch從.pt轉成了.onnx
1. 模型載入
使用TensorRt載入onnx模型的步驟其實是很固定的,根據官方例呈給出的示範,載入一個onnx的模型分為以下幾步
- 建立builder(構建器)
- 建立網路定義：builder —> network
- 配置引數：builder —> config
- 生成engine：builder —> engine (network, config)
- 序列化儲存：engine —> serialize
- 釋放資源：delete
- 第一步是使用TensorRT的api來宣告一個構建器型別。
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <logger.h>
#include <NvOnnxParser.h>
nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(sample::gLogger);
構建器初始化引數需要傳入一個gLogger物件，用於構建時的日誌儲存與列印。
- 接著使用構建器建立計算圖網路，也就是先建立一個空網路。建立時我們要指定顯式batchsize的大小，一般在部署環節的batchsize都設定為1。
const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
雖然第一行程式碼很長，但是實際上這只是TensorRT官方給出的一個靜態物件變數，值為0,因此explicitBatch相當於1 << 0 = 1,因此這裡相當於
const auto explicitBatch = 1U;
nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
- 接著我們來建立一個onnx的解析器來對onnx模型進行解析。
auto parser = nvonnxparser::createParser(*network, sample::gLogger);
parser->parseFromFile(onnx_file.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kINFO));
其中parserFromFile的第一個引數是讀取到的onnx檔案的地址，可以這樣獲取
string onnx_file = "./MNIST.onnx";
onnx_file.c_str();
- 接著我們來建立一個config物件，對網路的一些引數進行設定
auto config = builder->createBuilderConfig();
if (builder->platformHasFastFp16())
config->setFlag(nvinfer1::BuilderFlag::kFP16); // 若裝置支援FP16推理，則使用FP16模式
size_t free, total;
cudaMemGetInfo(&free, &total); // 獲取裝置視訊記憶體資訊
config->setMaxWorkspaceSize(free); // 將所有空餘視訊記憶體用於推理
其中必須設定的只有setMaxWorkspaceSize這一項。
為了在後面獲取輸入輸出時能知曉輸入輸出的shape，這裡為輸入輸出繫結名稱
network->getInput(0)->setName("input");
network->getOutput(0)->setName("output");
getInput的引數0代表輸入的張量索引，因為我們的輸入輸出都只有一個張量，因此填索引0即可。
- 最後建立TensorRT最關鍵的用於推理的Engine(引擎)
auto engine = builder->buildEngineWithConfig(*network, *config);
至此，構建階段就算完成了。至於engine的序列化為模型檔案以及透過讀取模型檔案來載入engine這裡線掠過。
- 最後，釋放除了engine之外的物件或者說指標。不難發現TensorRT中資料的傳遞都是使用指標而非引用來實現的，不知道為啥。
delete config;
delete parser;
delete network;
delete builder;
//or
config->destory();
parser->destory();
network->destory();
builder->destory();
2. 推理階段
在推理階段要做的工作依然很多，好在這次部署的模型較為簡單。推理階段的工作主要分為以下幾點。
- 對輸入的資料(影像)進行預處理操作，使之符合網路的輸入要求。
- 在GPU上申請需要的記憶體，並使用engine的context(上下文)成員進行推理
- 獲取輸出後對輸出進行後處理(argmax,nms等)，以獲取有用的資訊。
- 首先對於輸入的影像，並不是直接就可以送進網路進行推理的，因為其格式可能不符合網路的要求，因此需要對其做一些變換以適配網路的輸入格式。
可以使用Netron來檢視onnx模型的網路結構，瞭解輸入輸出的格式。
可以看出我要部署的網路的輸入格式為1×1×32×32(b, c, h, w)的一個張量，輸出為1×10的一個張量。因此對於輸入影像的格式，必須為單通道，大小為32×32畫素才可以輸入到網路中。
這裡選用opencv對影像進行讀取，並將影像透過resize，轉換為灰度圖(單通道)，然後進行歸一化，從而使之符合網路的輸入格式。
cv::Mat img = cv::imread("../five.png");
cv::cvtColor(img, img, cv::COLOR_BGR2GRAY);
cv::resize(img, img, {32, 32});
img = ~img;
img.convertTo(img, CV_32F); // 轉換為浮點型
img /= 255; // 歸一化
由於訓練時的MNIST資料集中的圖片為黑底白字的影像，而網圖多為白底黑字，因此使用img = ~img來進行反色。可以根據實際場景進行調整。
實際上opencv的圖片格式儲存是(h , w, c),也就是(高, 寬, 通道數)，而pytorch等主流框架都預設張量的輸入格式為(c,h,w),上面對模型的檢視也驗證了，這一點。因此我們還需要將opencv的(h,w,c)格式轉換為(c,h,w)格式。轉換程式碼如下。
// 原始影像，尺寸為(h, w, c)
void hwc2chw(cv::Mat &image)
{
int h = image.rows;
int w = image.cols;
int c = image.channels();
// 尺寸轉換為(h*w, c, 1)，此步驟不對記憶體進行修改
image = image.reshape(1, h * w);
// 影像轉置，尺寸變為(c, h*w, 1)
image = image.t();
// 尺寸轉換為(c, h, w)，此步驟不對記憶體進行修改
image = image.reshape(w, c);
}
引用於部落格
- 接下來就是在cuda上申請記憶體，並在其中進行推理。
在申請記憶體之前，顯然得先知道申請多少。再推理過程中，我們需要為輸入以及輸入各申請一塊記憶體，因此需要統計輸入輸出的元素數量即可。
/*獲取輸入輸出的idx進而獲取其維度(dims)*/
auto input_idx = engine->getBindingIndex("input");
auto output_idx = engine->getBindingIndex("output");
auto input_dims = engine->getBindingDimensions(input_idx);
auto output_dims = engine->getBindingDimensions(output_idx);
int input_sz = 1, output_sz = 1;
/*獲取需要申請存放輸入輸出的視訊記憶體大小，.nbDims返回維度數，d介面訪問每一維度的成員個數*/
for (int i = 0; i < input_dims.nbDims; i++)
input_sz *= input_dims.d[i];
for (int i = 0; i < output_dims.nbDims; i++)
output_sz *= output_dims.d[i];
/*這樣統計過後，input_sz和output_sz分別即為輸入輸出元素的大小，即需要申請記憶體的個數(float型)*/
寫完上面才知道shift+tab可以讓程式碼塊整體後退一個tab -_-
- 接著便是申請記憶體然後推理了。
申請記憶體
void *device_buffer[2]; // 裝置記憶體
float *output_buffer; // 輸出結果記憶體
cudaMalloc(&device_buffer[input_idx], input_sz * sizeof(float));
cudaMalloc(&device_buffer[output_idx], output_sz * sizeof(float));
output_buffer = new float[output_sz];
順便建立一個cuda流用於cuda自動管理非同步的記憶體管理
cudaStream_t stream;
cudaStreamCreate(&stream);
推理
auto context = engine->createExecutionContext(); // 建立context進行推理
cudaMemcpyAsync(device_buffer[input_idx], img.data, input_sz * sizeof(float), cudaMemcpyHostToDevice, stream); // 將預處理好的資料複製到cuda記憶體上
context->enqueueV2(device_buffer, stream, nullptr);// 加入推理佇列進行推理,推理結果也會存入device_buffer
cudaMemcpyAsync(output_buffer, device_buffer[output_idx], output_sz * sizeof(float), cudaMemcpyDeviceToHost, stream);// 推理結束後將結果複製到本地的output_buffer中
cudaStreamSynchronize(stream); // 等待流同步，即阻塞，直到以上所有的操作完成
對於cudaMemcpyAsync()，這是一個非同步的複製，會將資料複製到指定位置，決定是上傳還是下載的是第四個引數，cudaMemcpyHostToDevice代表本地資料上傳到cuda，而cudaMemcpyDeviceToHost代表cuda資料下載到本地。
- 最後便是對推理的資料進行後處理來獲取有效資訊。
上面透過使用Netron檢視了該模型的輸出張量維度為1×10,但是1×10張量的具體含義是訓練時我們知道的由資料集的標籤來指定的。因此在部署模型前一定要知道輸出的每個部分代表什麼意思，要不然就無法解析輸出。
對於這裡的MNIST資料集，其1×10的張量即0-9的陣列中儲存的是該張圖片為該數字的機率程度，越大說明越接近該數字。因此對於每個輸出，我們只需要找到1×10張量中最大的那個元素的下標，就是對應模型推理出的該圖片上的數字。這個按某個維度尋找最大值下標的操作就稱為argmax()。由於c++中我找不到現成的介面，就簡單寫了一個。
int argmax(float *output, int len)
{
float Max = -1;
int Max_idx = -1;
for (int i = 0; i < len; i++)
(output[i] > Max ? Max = output[i], Max_idx = i : false);
return Max_idx;
}
透過argmax操作獲取數字最大的成員的下標，即可認為該下標即為數字的型別。至此，整個模型的從構建到推理出結果就結束了。
最後別忘了將上面用到的記憶體以及engine釋放
delete[] output_buffer;
cudaFree(device_buffer[output_idx]);
cudaFree(device_buffer[input_idx]);
cudaStreamDestroy(stream);
delete engine;
or
engine->destory();
最後將整個程式碼貼上來
這裡是程式碼
#include <opencv4/opencv2/opencv.hpp>
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <logger.h>
#include "/home/ruby/Tensorrt/trt/include/NvInfer.h"
#include <NvOnnxParser.h>
#include <filesystem>
#include <fstream>
#include <iostream>
const std::string onnx_file = "/home/ruby/Desktop/work/CV/C++Test/MNIST_TEST.onnx";
int argmax(float *output, int len)
{
float Max = -1;
int Max_idx = -1;
for (int i = 0; i < len; i++)
(output[i] > Max ? Max = output[i], Max_idx = i : false);
return Max_idx;
}
int main()
{
/*1. 構建 */
std::filesystem::path onnx_file_path(onnx_file);
nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(sample::gLogger);
const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
auto parser = nvonnxparser::createParser(*network, sample::gLogger);
parser->parseFromFile(onnx_file.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kINFO));
network->getInput(0)->setName("input");
network->getOutput(0)->setName("output");
auto config = builder->createBuilderConfig();
if (builder->platformHasFastFp16())
config->setFlag(nvinfer1::BuilderFlag::kFP16);
size_t free, total;
cudaMemGetInfo(&free, &total);
config->setMaxWorkspaceSize(free);
auto engine = builder->buildEngineWithConfig(*network, *config);
delete config;
delete parser;
delete network;
delete builder;
/*2. 資料預處理*/
cv::Mat img = cv::imread("../five.png");
cv::cvtColor(img, img, cv::COLOR_BGR2GRAY);
cv::resize(img, img, {32, 32});
img = ~img;
img.convertTo(img, CV_32F);
img /= 255;
/*3. 申請視訊記憶體用於推理*/
auto input_idx = engine->getBindingIndex("input");
auto output_idx = engine->getBindingIndex("output");
auto input_dims = engine->getBindingDimensions(input_idx);
auto output_dims = engine->getBindingDimensions(output_idx);
int input_sz = 1, output_sz = 1;
/*獲取需要申請存放輸入輸出的視訊記憶體大小，nbDims返回維度數，d介面訪問每一維度的成員個數*/
for (int i = 0; i < input_dims.nbDims; i++)
input_sz *= input_dims.d[i];
for (int i = 0; i < output_dims.nbDims; i++)
output_sz *= output_dims.d[i];
void *device_buffer[2];
float *output_buffer;
cudaMalloc(&device_buffer[input_idx], input_sz * sizeof(float));
cudaMalloc(&device_buffer[output_idx], output_sz * sizeof(float));
cudaStream_t stream;
cudaStreamCreate(&stream);
output_buffer = new float[output_sz];
/*4. 開始推理*/
auto context = engine->createExecutionContext();
cudaMemcpyAsync(device_buffer[input_idx], img.data, input_sz * sizeof(float), cudaMemcpyHostToDevice, stream);
context->enqueueV2(device_buffer, stream, nullptr);
cudaMemcpyAsync(output_buffer, device_buffer[output_idx], output_sz * sizeof(float), cudaMemcpyDeviceToHost, stream);
cudaStreamSynchronize(stream);
/* 後處理 */
int label = argmax(output_buffer, output_sz);
std::cout << "預測為數字" << label << std::endl;
delete[] output_buffer;
cudaFree(device_buffer[output_idx]);
cudaFree(device_buffer[input_idx]);
cudaStreamDestroy(stream);
delete engine;
delete context;
return 0;
}
這裡插一嘴，如果想在markdown中插入程式碼塊的話，可以這樣寫
<details>
<summary> 我是摺疊標題 </summary>
<code>
//這裡有空格
```c++//在這裡開始
```//在這裡結束
</code>
</details>
補充,關於上交TRTModule
在上交開源的TRTModule.cpp中，也大致遵循了上述的流程來進行搭建和推理，但是由於裝甲板檢測使用的YOLO模型(好像是YOLOFACE)，因此對圖片的預處理以及後處理會有所不同，而且後處理操作由於要實現非極大值抑制，比較麻煩。
還有一點就是上交在網路的輸出後又加了幾個層來對輸出進行在網路層面的後處理。
auto yolov5_output = network->getOutput(0);
/*具體目標為將輸出的置信度部分提取，並提取最大的topk個數作為輸出*/
auto slice_layer = network->addSlice(*yolov5_output, Dims3{0, 0, 8}, Dims3{1, 15120, 1}, Dims3{1, 1, 1});
auto yolov5_conf = slice_layer->getOutput(0);
/*reshape，(1,15120,1)->(1,15120)*/
auto shuffle_layer = network->addShuffle(*yolov5_conf);
shuffle_layer->setReshapeDimensions(Dims2{1, 15120});
yolov5_conf = shuffle_layer->getOutput(0);
/*topk，提取最大的前topk個元素*/
auto topk_layer = network->addTopK(*yolov5_conf, TopKOperation::kMAX, TOPK_NUM, 1 << 1);
auto topk_idx = topk_layer->getOutput(1);
/*透過topk層的索引來重塑張量*/
auto gather_layer = network->addGather(*yolov5_output, *topk_idx, 1);
gather_layer->setNbElementWiseDims(1);
auto yolov5_output_topk = gather_layer->getOutput(0);
/*繫結輸入輸出，防止被最佳化掉*/
yolov5_output_topk->setName("output-topk");
network->getInput(0)->setName("input");
/*繫結新輸出*/
network->markOutput(*yolov5_output_topk);
/*解綁舊輸出，解綁的張量會被當作暫時量被最佳化掉*/
network->unmarkOutput(*yolov5_output);
這裡講幾個點。這幾步找了好幾天的資料才看懂。
第一步銜接的是剛構造完計算圖(network)，我們取獲取網路的輸出張量。然後接著在網路的後面新增了各種層
addslice，在網路後新增一個切片層，輸入為上一步的輸出也就是yolov5_output。引數傳入的是三個維度，分別為起始位置，切片完後的shape，以及步長。這裡意味著對於每個張量從第8位也就是第九個資料進行切片，步長為1。值得一提的是在上交的onnx模型中，第八位的輸出是該張圖片作為一個裝甲板的置信度。因此此步操作就是將(1,15120,20)的張量中關於置信度的那一維度切片出來成為(1,15120,1)的張量
addshuffle，在網路後再新增一個shuffle層，相當於網路中做了一步reshape操作，從(1,15120,1)輸出(1,15120)的二維張量。
addTopk，尋找沿著張量的某個維度的滿足最條件的一些量。這裡的傳參使用
TopKOperation::kMAX也就是選擇最大的前topk個元素，然後TOPK_NUM就是目標要找的前TOPK個數的多少，而最後一位，在註釋中叫做reduceAxes，使用很奇怪的方式來指定在那個維度上進行topk操作。註釋顯示使用位掩碼來指定，即當最後一個引數傳入5(101)時，僅有第0維和第2維進行操作，而第1維不操作。因此這裡傳入1<<1(10)代表這僅對第1維進行操作，而不對第0維進行操作。這十分合理，因為對於(1×15120)的二維張量，第0維每個元素僅有一個元素，topk是沒有意義的，而對於第1維的topk可以排出置信度前TOP_NUM大的樣本。
接著對於topk層的輸出，一個是第0維的輸出
getOutput(0),輸出的是前topk的值，而第1維的輸出getOuput(1),輸出的是前topk的輸出的索引位置的張量。
這裡只關心位置。因為我們提取置信度為前topk的張量的目的就是獲取其索引然後透過索引去從(1,15120,20)->(1,NUM_TOPK,20)。從而減少要處理的樣本量。而這透過索引重建資料的操作就交給了最後一個gather層。gather層透過提供的索引來聚集對應的張量。這樣就實現了從索引到(1,NUM_TOPK,20)的實現。
最後將原先的輸出解綁，繫結從最後的gather層的輸出為整個網路的輸出。
最後貼上上交的後處理過程，僅有註釋，今天有點豬腦過載就不寫解析了。不過還是值得提一嘴上交的輸出20個元素中，0-7為四點，左上角開始/逆時針(x,y)，8為置信度，9-12為四種顏色的可能(r,b,g,p),13-19為七種裝甲板型別id。
// post-process [nms]
std::vector<bbox_t> rst;
rst.reserve(TOPK_NUM);
std::vector<uint8_t> removed(TOPK_NUM);
auto input_dims = engine->getBindingDimensions(input_idx);
auto output_dims = engine->getBindingDimensions(output_idx);
for (int i = 0; i < TOPK_NUM; i++)
{
auto *box_buffer = output_buffer + i * 20; // 20->23
/*第8位推斷為是<裝甲板>的置信度(未經過sigmoid歸一化)，而keep_thres為置信閾值，透過sigmoid的反函式來求出未經過sigmoid的置信度*/
/*兩者相比較，篩掉置信度低於置信閾值的樣本*/
if (box_buffer[8] < inv_sigmoid(KEEP_THRES))
break;
/*判斷*/
if (removed[i]) // 只處理沒被romove的樣本
continue;
/*向rst中填入一個空成員*/
rst.emplace_back();
/*取出最後一個成員，也就是最後一個空成員*/
auto &box = rst.back();
/*將本輪迴圈樣本的前8位資料傳給box的pts，即畫素四點，每點(x,y),四點為 4 * 2 = 8 個資料*/
memcpy(&box.pts, box_buffer, 8 * sizeof(float));
/*將四點按開頭算的比例對映到原影像上*/
for (auto &pt : box.pts)
pt.x *= fx, pt.y *= fy;
/*讀取buffer不同的資料段透過argmax來獲取該裝甲板的資訊*/
box.confidence = sigmoid(box_buffer[8]);
box.color_id = argmax(box_buffer + 9, 4);
box.tag_id = argmax(box_buffer + 13, 7);
/*透過計算IOU來進行非極大值抑制*/
/*
* 可以注意到，由於我們檢測裝甲板的任務很簡單，目標少，因此直接採用有交集就篩掉的原則，不再用傳統NMS那一套
* 比如先對置信度排序，然後遍歷該種類的所有box，計算IOU，篩掉IOU大於閾值或者置信度小於閾值的box
*/
for (int j = i + 1; j < TOPK_NUM; j++) // 遍歷每一個樣本
{
auto *box2_buffer = output_buffer + j * 20;
if (box2_buffer[8] < inv_sigmoid(KEEP_THRES)) // 過掉置信度小的樣本
break;
if (removed[j])
continue;
if (is_overlap(box_buffer, box2_buffer)) // 如果有交集，直接將其remove
removed[j] = true;
}
}
return rst;