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，大小为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 推理过程

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 ---> 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++)基础代码解析