TensorRT(C++)基础代码解析
前言
一、TensorRT工作流程
二、C++ API
2.1 构建阶段
TensorRT build engine的过程
- 创建builder
- 创建网络定义：builder —> network
- 配置参数：builder —> config
- 生成engine：builder —> engine (network, config)
- 序列化保存：engine —> serialize
- 释放资源：delete
2.1.1 创建builder
nvinfer1 是 NVIDIA TensorRT 的 C++ 接口命名空间。构建阶段的最高级别接口是 Builder。Builder负责优化一个模型，并产生Engine。通过如下接口创建一个Builder。
nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);
2.1.2 创建网络定义
NetworkDefinition接口被用来定义模型。接口createNetworkV2接受配置参数，参数用按位标记的方式传入。比如上面激活explicitBatch，是通过1U << static_cast<uint32_t (nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); 将explicitBatch对应的配置位设置为1实现的。在新版本中，请使用createNetworkV2而非其他任何创建NetworkDefinition 的接口。
auto explicitBatch = 1U << static_cast<uint32_t
(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
// 调用createNetworkV2创建网络定义，参数是显性batch
nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
2.1.3 定义网络结构
将模型转移到TensorRT的最常见的方式是以ONNX格式从框架中导出（将在后续课程进行介绍），并使用TensorRT的ONNX解析器来填充网络定义。同时，也可以使用TensorRT的Layer和Tensor等接口一步一步地进行定义。通过接口来定义网络的代码示例如下：
添加输入层
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 定义网络输入输出
定义哪些张量是网络的输入和输出。没有被标记为输出的张量被认为是瞬时值，可以被构建者优化掉。输入和输出张量必须被命名，以便在运行时，TensorRT知道如何将输入和输出缓冲区绑定到模型上。
// 设置输出名字
sigmoid->getOutput(0)->setName("output");
// 标记输出，没有标记会被当成顺时针优化掉
network->markOutput(*sigmoid->getOutput(0));
2.1.5 配置参数
添加相关Builder 的配置。createBuilderConfig接口被用来指定TensorRT应该如何优化模型
nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
// 设置最大工作空间大小，单位是字节
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. 创建builder
8. 创建网络定义：builder ---> network
9. 配置参数：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;
// 保存权重
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)); // 保存权重的大小
outfile.write((char *)(data), size * sizeof(float)); // 保存权重的数据
outfile.close();
}
// 读取权重
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)); // 读取权重的大小
std::vector<float> data(size); // 创建一个vector，大小为size
infile.read((char *)(data.data()), size * sizeof(float)); // 读取权重的数据
infile.close();
return data;
}
int main()
{
// ======= 1. 创建builder =======
TRTLogger logger;
nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);
// ======= 2. 创建网络定义：builder ---> network =======
// 显性batch
// 1 << 0 = 1，二进制移位，左移0位，相当于1（y左移x位，相当于y乘以2的x次方）
auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
// 调用createNetworkV2创建网络定义，参数是显性batch
nvinfer1::INetworkDefinition *network = builder->createNetworkV2(explicitBatch);
// 定义网络结构
// mlp多层感知机：input(1,3,1,1) --> fc1 --> sigmoid --> output (2)
// 创建一个input tensor ，参数分别是：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});
// 创建全连接层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};
// 将权重保存到文件中，演示从别的来源加载权重
saveWeights("model/fc1.wts", fc1_weight_data, 6);
saveWeights("model/fc1.bias", fc1_bias_data, 2);
// 读取权重
auto fc1_weight_vec = loadWeights("model/fc1.wts");
auto fc1_bias_vec = loadWeights("model/fc1.bias");
// 转为nvinfer1::Weights类型，参数分别是：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;
// 调用addFullyConnected创建全连接层，参数分别是：input tensor, output size, weight, bias
nvinfer1::IFullyConnectedLayer *fc1 = network->addFullyConnected(*input, output_size, fc1_weight, fc1_bias);
// 添加sigmoid激活层，参数分别是：input tensor, activation type（激活函数类型）
nvinfer1::IActivationLayer *sigmoid = network->addActivation(*fc1->getOutput(0), nvinfer1::ActivationType::kSIGMOID);
// 设置输出名字
sigmoid->getOutput(0)->setName("output");
// 标记输出，没有标记会被当成顺时针优化掉
network->markOutput(*sigmoid->getOutput(0));
// 设定最大batch size
builder->setMaxBatchSize(1);
// ====== 3. 配置参数：builder ---> config ======
// 添加配置参数，告诉TensorRT应该如何优化网络
nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
// 设置最大工作空间大小，单位是字节
config->setMaxWorkspaceSize(1 << 28); // 256MiB
// ====== 4. 创建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. 释放资源 ======
// 理论上，这些资源都会在程序结束时自动释放，但是为了演示，这里手动释放部分
outfile.close();
delete serialized_engine;
delete engine;
delete config;
delete network;
delete builder;
std::cout << "engine文件生成成功！" << std::endl;
return 0;
}
2.2 运行期
TensorRT runtime 推理过程
- 创建一个runtime对象
- 反序列化生成engine：runtime —> engine
- 创建一个执行上下文ExecutionContext：engine —> context
- 填充数据
- 执行推理：context —> enqueueV2
- 释放资源：delete
2.2.1 创建一个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 创建一个执行上下文ExecutionContext
从Engine创建的ExecutionContext接口是调用推理的主要接口。ExecutionContext包含与特定调用相关的所有状态，因此可以有多个与单个引擎相关的上下文，且并行运行它们。
nvinfer1::IExecutionContext *context = engine->createExecutionContext();
2.2.4 为推理填充输入
首先创建CUDA Stream用于推理的执行。
cudaStream_t stream = nullptr;
cudaStreamCreate(&stream);
同时在CPU和GPU上分配输入输出内存，并将输入数据从CPU拷贝到GPU上。
// 输入数据
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;
// 输出数据
float* h_out_data = new float[2]{0.0, 0.0};
int out_data_size = sizeof(float) * 2;
float* d_out_data = nullptr;
// 申请GPU上的内存
cudaMalloc(&d_in_data, in_data_size);
cudaMalloc(&d_out_data, out_data_size);
// 拷贝数据
cudaMemcpyAsync(d_in_data, h_in_data, in_data_size, cudaMemcpyHostToDevice, stream);
// enqueueV2中是把输入输出的内存地址放到bindings这个数组中，需要写代码时确定这些输入输出的顺序（这样容易出错，而且不好定位bug，所以新的接口取消了这样的方式，不过目前很多官方 sample 也在用v2）
float* bindings[] = {d_in_data, d_out_data};
2.2.4 调用enqueueV2来执行推理
bool success = context -> enqueueV2((void **) bindings, stream, nullptr);
// 数据从device --> host
cudaMemcpyAsync(host_output_data, device_output_data, output_data_size, cudaMemcpyDeviceToHost, stream);
// 等待流执行完毕
cudaStreamSynchronize(stream);
// 输出结果
std::cout << "输出结果: " << 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;
完整代码文章来源：https://www.toymoban.com/news/detail-799934.html
/*
使用.cu是希望使用CUDA的编译器NVCC，会自动连接cuda库
TensorRT runtime 推理过程
1. 创建一个runtime对象
2. 反序列化生成engine：runtime ---> engine
3. 创建一个执行上下文ExecutionContext：engine ---> context
4. 填充数据
5. 执行推理：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); // 创建一个vector，大小为size
file.seekg(0, std::ios::beg); // 定位到文件开头
file.read((char *)data.data(), size); // 读取文件内容到data中
file.close();
return data;
}
int main()
{
// ==================== 1. 创建一个runtime对象 ====================
TRTLogger logger;
nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);
// ==================== 2. 反序列化生成engine ====================
// 读取文件
auto engineModel = loadEngineModel("./model/mlp.engine");
// 调用runtime的反序列化方法，生成engine，参数分别是：模型数据地址，模型大小，pluginFactory
nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(engineModel.data(), engineModel.size(), nullptr);
if (!engine)
{
std::cout << "deserialize engine failed!" << std::endl;
return -1;
}
// ==================== 3. 创建一个执行上下文 ====================
nvinfer1::IExecutionContext *context = engine->createExecutionContext();
// ==================== 4. 填充数据 ====================
// 设置stream 流
cudaStream_t stream = nullptr;
cudaStreamCreate(&stream);
// 数据流转：host --> device ---> inference ---> host
// 输入数据
float *host_input_data = new float[3]{2, 4, 8}; // host 输入数据
int input_data_size = 3 * sizeof(float); // 输入数据大小
float *device_input_data = nullptr; // device 输入数据
// 输出数据
float *host_output_data = new float[2]{0, 0}; // host 输出数据
int output_data_size = 2 * sizeof(float); // 输出数据大小
float *device_output_data = nullptr; // device 输出数据
// 申请device内存
cudaMalloc((void **)&device_input_data, input_data_size);
cudaMalloc((void **)&device_output_data, output_data_size);
// host --> device
// 参数分别是：目标地址，源地址，数据大小，拷贝方向
cudaMemcpyAsync(device_input_data, host_input_data, input_data_size, cudaMemcpyHostToDevice, stream);
// bindings告诉Context输入输出数据的位置
float *bindings[] = {device_input_data, device_output_data};
// ==================== 5. 执行推理 ====================
bool success = context -> enqueueV2((void **) bindings, stream, nullptr);
// 数据从device --> host
cudaMemcpyAsync(host_output_data, device_output_data, output_data_size, cudaMemcpyDeviceToHost, stream);
// 等待流执行完毕
cudaStreamSynchronize(stream);
// 输出结果
std::cout << "输出结果: " << 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;
}
总结
TensorRT(C++)基础代码解析文章来源地址https://www.toymoban.com/news/detail-799934.html
到了这里，关于TensorRT(C++)基础代码解析的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！
