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

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


文章目录

  • TensorRT(C++)基础代码解析
  • 前言
  • 一、TensorRT工作流程
  • 二、C++ API
    • 2.1 构建阶段
      • 2.1.1 创建builder
      • 2.1.2 创建网络定义
      • 2.1.3 定义网络结构
      • 2.1.4 定义网络输入输出
      • 2.1.5 配置参数
      • 2.1.6 生成Engine
      • 2.1.7 保存为模型文件
      • 2.1.8 释放资源
    • 2.2 运行期
      • 2.2.1 创建一个runtime对象
      • 2.2.2 反序列化生成engine
      • 2.2.3 创建一个执行上下文ExecutionContext
      • 2.2.4 为推理填充输入
      • 2.2.4 调用enqueueV2来执行推理
      • 2.2.5 释放资源
  • 总结


前言


一、TensorRT工作流程

TensorRT(C++)基础代码解析_第1张图片

TensorRT(C++)基础代码解析_第2张图片

二、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

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 
#include 
#include 
#include 

#include 

// 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(C++)基础代码解析_第3张图片
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 
#include 
#include 
#include 

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

你可能感兴趣的:(c++,开发语言)