TensorRT(C++)基础代码解析

TensorRT(C++)基础代码解析

---

---

前言

---

一、TensorRT工作流程

二、C++ API

2.1 构建阶段

TensorRT build engine的过程

1. 创建builder
2. 创建网络定义：builder —> network
3. 配置参数：builder —> config
4. 生成engine：builder —> engine (network, config)
5. 序列化保存：engine —> serialize
6. 释放资源：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，大小为sizeinfile.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创建网络定义，参数是显性batchnvinfer1::INetworkDefinition *network = builder->createNetworkV2(explicitBatch);// 定义网络结构// mlp多层感知机：input(1,3,1,1) --> fc1 --> sigmoid --> output (2)// 创建一个input tensor ，参数分别是：name, data type, dimsconst int input_size = 3;nvinfer1::ITensor *input = network->addInput("data", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, input_size, 1, 1});// 创建全连接层fc1// weight and biasconst 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, sizenvinfer1::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, biasnvinfer1::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 sizebuilder->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 推理过程

1. 创建一个runtime对象
2. 反序列化生成engine：runtime —> engine
3. 创建一个执行上下文ExecutionContext：engine —> context
4. 填充数据
5. 执行推理：context —> enqueueV2
6. 释放资源：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;

完整代码

/*
使用.cu是希望使用CUDA的编译器NVCC，会自动连接cuda库TensorRT runtime 推理过程1. 创建一个runtime对象
2. 反序列化生成engine：runtime ---> engine
3. 创建一个执行上下文ExecutionContext：engine ---> context4. 填充数据5. 执行推理：context ---> enqueueV26. 释放资源：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，大小为sizefile.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，参数分别是：模型数据地址，模型大小，pluginFactorynvinfer1::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 --> hostcudaMemcpyAsync(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++)基础代码解析