TensorRT(C++)基礎(chǔ)代碼解析
前言
一、TensorRT工作流程
二、C++ API
2.1 構(gòu)建階段
TensorRT build engine的過程
- 創(chuàng)建builder
- 創(chuàng)建網(wǎng)絡(luò)定義：builder —> network
- 配置參數(shù)：builder —> config
- 生成engine：builder —> engine (network, config)
- 序列化保存：engine —> serialize
- 釋放資源：delete
2.1.1 創(chuàng)建builder
nvinfer1 是 NVIDIA TensorRT 的 C++ 接口命名空間。構(gòu)建階段的最高級別接口是 Builder。Builder負責優(yōu)化一個模型，并產(chǎn)生Engine。通過如下接口創(chuàng)建一個Builder。
nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);
2.1.2 創(chuàng)建網(wǎng)絡(luò)定義
NetworkDefinition接口被用來定義模型。接口createNetworkV2接受配置參數(shù)，參數(shù)用按位標記的方式傳入。比如上面激活explicitBatch，是通過1U << static_cast<uint32_t (nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); 將explicitBatch對應(yīng)的配置位設(shè)置為1實現(xiàn)的。在新版本中，請使用createNetworkV2而非其他任何創(chuàng)建NetworkDefinition 的接口。
auto explicitBatch = 1U << static_cast<uint32_t
(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
// 調(diào)用createNetworkV2創(chuàng)建網(wǎng)絡(luò)定義，參數(shù)是顯性batch
nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
2.1.3 定義網(wǎng)絡(luò)結(jié)構(gòu)
將模型轉(zhuǎn)移到TensorRT的最常見的方式是以O(shè)NNX格式從框架中導出（將在后續(xù)課程進行介紹），并使用TensorRT的ONNX解析器來填充網(wǎng)絡(luò)定義。同時，也可以使用TensorRT的Layer和Tensor等接口一步一步地進行定義。通過接口來定義網(wǎng)絡(luò)的代碼示例如下：
添加輸入層
const int input_size = 3;
nvinfer1::ITensor *input = network->addInput("data", nvinfer1::DataType::kFLOAT,nvinfer1::Dims4{1, input_size, 1, 1})
添加全連接層
nvinfer1::IFullyConnectedLayer* fc1 = network->addFullyConnected(*input, output_size, fc1w, fc1b);
添加激活層
nvinfer1::IActivationLayer* relu1 = network->addActivation(*fc1->getOutput(0), nvinfer1::ActivationType::kRELU);
2.1.4 定義網(wǎng)絡(luò)輸入輸出
定義哪些張量是網(wǎng)絡(luò)的輸入和輸出。沒有被標記為輸出的張量被認為是瞬時值，可以被構(gòu)建者優(yōu)化掉。輸入和輸出張量必須被命名，以便在運行時，TensorRT知道如何將輸入和輸出緩沖區(qū)綁定到模型上。
// 設(shè)置輸出名字
sigmoid->getOutput(0)->setName("output");
// 標記輸出，沒有標記會被當成順時針優(yōu)化掉
network->markOutput(*sigmoid->getOutput(0));
2.1.5 配置參數(shù)
添加相關(guān)Builder 的配置。createBuilderConfig接口被用來指定TensorRT應(yīng)該如何優(yōu)化模型
nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
// 設(shè)置最大工作空間大小，單位是字節(jié)
config->setMaxWorkspaceSize(1 << 28); // 256MiB
2.1.6 生成Engine
nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
2.1.7 保存為模型文件
nvinfer1::IHostMemory *serialized_engine = engine->serialize();
// 存入文件
std::ofstream outfile("model/mlp.engine", std::ios::binary);
assert(outfile.is_open() && "Failed to open file for writing");
outfile.write((char *)serialized_engine->data(), serialized_engine->size());
2.1.8 釋放資源
outfile.close();
delete serialized_engine;
delete engine;
delete config;
delete network;
完整代碼
/*
TensorRT build engine的過程
7. 創(chuàng)建builder
8. 創(chuàng)建網(wǎng)絡(luò)定義：builder ---> network
9. 配置參數(shù)：builder ---> config
10. 生成engine：builder ---> engine (network, config)
11. 序列化保存：engine ---> serialize
12. 釋放資源：delete
*/
#include <iostream>
#include <fstream>
#include <cassert>
#include <vector>
#include <NvInfer.h>
// logger用來管控打印日志級別
// TRTLogger繼承自nvinfer1::ILogger
class TRTLogger : public nvinfer1::ILogger
{
void log(Severity severity, const char *msg) noexcept override
{
// 屏蔽INFO級別的日志
if (severity != Severity::kINFO)
std::cout << msg << std::endl;
}
} gLogger;
// 保存權(quán)重
void saveWeights(const std::string &filename, const float *data, int size)
{
std::ofstream outfile(filename, std::ios::binary);
assert(outfile.is_open() && "save weights failed"); // assert斷言，如果條件不滿足，就會報錯
outfile.write((char *)(&size), sizeof(int)); // 保存權(quán)重的大小
outfile.write((char *)(data), size * sizeof(float)); // 保存權(quán)重的數(shù)據(jù)
outfile.close();
}
// 讀取權(quán)重
std::vector<float> loadWeights(const std::string &filename)
{
std::ifstream infile(filename, std::ios::binary);
assert(infile.is_open() && "load weights failed");
int size;
infile.read((char *)(&size), sizeof(int)); // 讀取權(quán)重的大小
std::vector<float> data(size); // 創(chuàng)建一個vector，大小為size
infile.read((char *)(data.data()), size * sizeof(float)); // 讀取權(quán)重的數(shù)據(jù)
infile.close();
return data;
}
int main()
{
// ======= 1. 創(chuàng)建builder =======
TRTLogger logger;
nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);
// ======= 2. 創(chuàng)建網(wǎng)絡(luò)定義：builder ---> network =======
// 顯性batch
// 1 << 0 = 1，二進制移位，左移0位，相當于1（y左移x位，相當于y乘以2的x次方）
auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
// 調(diào)用createNetworkV2創(chuàng)建網(wǎng)絡(luò)定義，參數(shù)是顯性batch
nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
// 定義網(wǎng)絡(luò)結(jié)構(gòu)
// mlp多層感知機：input(1,3,1,1) --> fc1 --> sigmoid --> output (2)
// 創(chuàng)建一個input tensor ，參數(shù)分別是：name, data type, dims
const int input_size = 3;
nvinfer1::ITensor *input = network->addInput("data", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, input_size, 1, 1});
// 創(chuàng)建全連接層fc1
// weight and bias
const float *fc1_weight_data = new float[input_size * 2]{0.1, 0.2, 0.3, 0.4, 0.5, 0.6};
const float *fc1_bias_data = new float[2]{0.1, 0.5};
// 將權(quán)重保存到文件中，演示從別的來源加載權(quán)重
saveWeights("model/fc1.wts", fc1_weight_data, 6);
saveWeights("model/fc1.bias", fc1_bias_data, 2);
// 讀取權(quán)重
auto fc1_weight_vec = loadWeights("model/fc1.wts");
auto fc1_bias_vec = loadWeights("model/fc1.bias");
// 轉(zhuǎn)為nvinfer1::Weights類型，參數(shù)分別是：data type, data, size
nvinfer1::Weights fc1_weight{nvinfer1::DataType::kFLOAT, fc1_weight_vec.data(), fc1_weight_vec.size()};
nvinfer1::Weights fc1_bias{nvinfer1::DataType::kFLOAT, fc1_bias_vec.data(), fc1_bias_vec.size()};
const int output_size = 2;
// 調(diào)用addFullyConnected創(chuàng)建全連接層，參數(shù)分別是：input tensor, output size, weight, bias
nvinfer1::IFullyConnectedLayer *fc1 = network->addFullyConnected(*input, output_size, fc1_weight, fc1_bias);
// 添加sigmoid激活層，參數(shù)分別是：input tensor, activation type（激活函數(shù)類型）
nvinfer1::IActivationLayer *sigmoid = network->addActivation(*fc1->getOutput(0), nvinfer1::ActivationType::kSIGMOID);
// 設(shè)置輸出名字
sigmoid->getOutput(0)->setName("output");
// 標記輸出，沒有標記會被當成順時針優(yōu)化掉
network->markOutput(*sigmoid->getOutput(0));
// 設(shè)定最大batch size
builder->setMaxBatchSize(1);
// ====== 3. 配置參數(shù)：builder ---> config ======
// 添加配置參數(shù)，告訴TensorRT應(yīng)該如何優(yōu)化網(wǎng)絡(luò)
nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
// 設(shè)置最大工作空間大小，單位是字節(jié)
config->setMaxWorkspaceSize(1 << 28); // 256MiB
// ====== 4. 創(chuàng)建engine：builder ---> network ---> config ======
nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
if (!engine)
{
std::cerr << "Failed to create engine!" << std::endl;
return -1;
}
// ====== 5. 序列化engine ======
nvinfer1::IHostMemory *serialized_engine = engine->serialize();
// 存入文件
std::ofstream outfile("model/mlp.engine", std::ios::binary);
assert(outfile.is_open() && "Failed to open file for writing");
outfile.write((char *)serialized_engine->data(), serialized_engine->size());
// ====== 6. 釋放資源 ======
// 理論上，這些資源都會在程序結(jié)束時自動釋放，但是為了演示，這里手動釋放部分
outfile.close();
delete serialized_engine;
delete engine;
delete config;
delete network;
delete builder;
std::cout << "engine文件生成成功！" << std::endl;
return 0;
}
2.2 運行期
TensorRT runtime 推理過程
- 創(chuàng)建一個runtime對象
- 反序列化生成engine：runtime —> engine
- 創(chuàng)建一個執(zhí)行上下文ExecutionContext：engine —> context
- 填充數(shù)據(jù)
- 執(zhí)行推理：context —> enqueueV2
- 釋放資源：delete
2.2.1 創(chuàng)建一個runtime對象
TensorRT運行時的最高層級接口是Runtime
nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);
2.2.2 反序列化生成engine
通過讀取模型文件并反序列化，我們可以利用runtime生成Engine。
nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size(), nullptr);
2.2.3 創(chuàng)建一個執(zhí)行上下文ExecutionContext
從Engine創(chuàng)建的ExecutionContext接口是調(diào)用推理的主要接口。ExecutionContext包含與特定調(diào)用相關(guān)的所有狀態(tài)，因此可以有多個與單個引擎相關(guān)的上下文，且并行運行它們。
nvinfer1::IExecutionContext *context = engine->createExecutionContext();
2.2.4 為推理填充輸入
首先創(chuàng)建CUDA Stream用于推理的執(zhí)行。
cudaStream_t stream = nullptr;
cudaStreamCreate(&stream);
同時在CPU和GPU上分配輸入輸出內(nèi)存，并將輸入數(shù)據(jù)從CPU拷貝到GPU上。
// 輸入數(shù)據(jù)
float* h_in_data = new float[3]{1.4, 3.2, 1.1};
int in_data_size = sizeof(float) * 3;
float* d_in_data = nullptr;
// 輸出數(shù)據(jù)
float* h_out_data = new float[2]{0.0, 0.0};
int out_data_size = sizeof(float) * 2;
float* d_out_data = nullptr;
// 申請GPU上的內(nèi)存
cudaMalloc(&d_in_data, in_data_size);
cudaMalloc(&d_out_data, out_data_size);
// 拷貝數(shù)據(jù)
cudaMemcpyAsync(d_in_data, h_in_data, in_data_size, cudaMemcpyHostToDevice, stream);
// enqueueV2中是把輸入輸出的內(nèi)存地址放到bindings這個數(shù)組中，需要寫代碼時確定這些輸入輸出的順序（這樣容易出錯，而且不好定位bug，所以新的接口取消了這樣的方式，不過目前很多官方 sample 也在用v2）
float* bindings[] = {d_in_data, d_out_data};
2.2.4 調(diào)用enqueueV2來執(zhí)行推理
bool success = context -> enqueueV2((void **) bindings, stream, nullptr);
// 數(shù)據(jù)從device --> host
cudaMemcpyAsync(host_output_data, device_output_data, output_data_size, cudaMemcpyDeviceToHost, stream);
// 等待流執(zhí)行完畢
cudaStreamSynchronize(stream);
// 輸出結(jié)果
std::cout << "輸出結(jié)果: " << host_output_data[0] << " " << host_output_data[1] << std::endl;
2.2.5 釋放資源
cudaStreamDestroy(stream);
cudaFree(device_input_data_address);
cudaFree(device_output_data_address);
delete[] host_input_data;
delete[] host_output_data;
delete context;
delete engine;
delete runtime;
完整代碼文章來源：http://www.zghlxwxcb.cn/news/detail-799934.html
/*
使用.cu是希望使用CUDA的編譯器NVCC，會自動連接cuda庫
TensorRT runtime 推理過程
1. 創(chuàng)建一個runtime對象
2. 反序列化生成engine：runtime ---> engine
3. 創(chuàng)建一個執(zhí)行上下文ExecutionContext：engine ---> context
4. 填充數(shù)據(jù)
5. 執(zhí)行推理：context ---> enqueueV2
6. 釋放資源：delete
*/
#include <iostream>
#include <vector>
#include <fstream>
#include <cassert>
#include "cuda_runtime.h"
#include "NvInfer.h"
// logger用來管控打印日志級別
// TRTLogger繼承自nvinfer1::ILogger
class TRTLogger : public nvinfer1::ILogger
{
void log(Severity severity, const char *msg) noexcept override
{
// 屏蔽INFO級別的日志
if (severity != Severity::kINFO)
std::cout << msg << std::endl;
}
} gLogger;
// 加載模型
std::vector<unsigned char> loadEngineModel(const std::string &fileName)
{
std::ifstream file(fileName, std::ios::binary); // 以二進制方式讀取
assert(file.is_open() && "load engine model failed!"); // 斷言
file.seekg(0, std::ios::end); // 定位到文件末尾
size_t size = file.tellg(); // 獲取文件大小
std::vector<unsigned char> data(size); // 創(chuàng)建一個vector，大小為size
file.seekg(0, std::ios::beg); // 定位到文件開頭
file.read((char *)data.data(), size); // 讀取文件內(nèi)容到data中
file.close();
return data;
}
int main()
{
// ==================== 1. 創(chuàng)建一個runtime對象 ====================
TRTLogger logger;
nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);
// ==================== 2. 反序列化生成engine ====================
// 讀取文件
auto engineModel = loadEngineModel("./model/mlp.engine");
// 調(diào)用runtime的反序列化方法，生成engine，參數(shù)分別是：模型數(shù)據(jù)地址，模型大小，pluginFactory
nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(engineModel.data(), engineModel.size(), nullptr);
if (!engine)
{
std::cout << "deserialize engine failed!" << std::endl;
return -1;
}
// ==================== 3. 創(chuàng)建一個執(zhí)行上下文 ====================
nvinfer1::IExecutionContext *context = engine->createExecutionContext();
// ==================== 4. 填充數(shù)據(jù) ====================
// 設(shè)置stream 流
cudaStream_t stream = nullptr;
cudaStreamCreate(&stream);
// 數(shù)據(jù)流轉(zhuǎn)：host --> device ---> inference ---> host
// 輸入數(shù)據(jù)
float *host_input_data = new float[3]{2, 4, 8}; // host 輸入數(shù)據(jù)
int input_data_size = 3 * sizeof(float); // 輸入數(shù)據(jù)大小
float *device_input_data = nullptr; // device 輸入數(shù)據(jù)
// 輸出數(shù)據(jù)
float *host_output_data = new float[2]{0, 0}; // host 輸出數(shù)據(jù)
int output_data_size = 2 * sizeof(float); // 輸出數(shù)據(jù)大小
float *device_output_data = nullptr; // device 輸出數(shù)據(jù)
// 申請device內(nèi)存
cudaMalloc((void **)&device_input_data, input_data_size);
cudaMalloc((void **)&device_output_data, output_data_size);
// host --> device
// 參數(shù)分別是：目標地址，源地址，數(shù)據(jù)大小，拷貝方向
cudaMemcpyAsync(device_input_data, host_input_data, input_data_size, cudaMemcpyHostToDevice, stream);
// bindings告訴Context輸入輸出數(shù)據(jù)的位置
float *bindings[] = {device_input_data, device_output_data};
// ==================== 5. 執(zhí)行推理 ====================
bool success = context -> enqueueV2((void **) bindings, stream, nullptr);
// 數(shù)據(jù)從device --> host
cudaMemcpyAsync(host_output_data, device_output_data, output_data_size, cudaMemcpyDeviceToHost, stream);
// 等待流執(zhí)行完畢
cudaStreamSynchronize(stream);
// 輸出結(jié)果
std::cout << "輸出結(jié)果: " << host_output_data[0] << " " << host_output_data[1] << std::endl;
// ==================== 6. 釋放資源 ====================
cudaStreamDestroy(stream);
cudaFree(device_input_data);
cudaFree(device_output_data);
delete host_input_data;
delete host_output_data;
delete context;
delete engine;
delete runtime;
return 0;
}
總結(jié)
TensorRT(C++)基礎(chǔ)代碼解析文章來源地址http://www.zghlxwxcb.cn/news/detail-799934.html
到了這里，關(guān)于TensorRT(C++)基礎(chǔ)代碼解析的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
