项目四：TensorRT编程

TensorRT编程 - 单层感知机结构

build.cpp

/*
一般先是运行 build.cpp, 在运行 runtime.cu
TensorRT命名空间是 nvinfer1;
1.创建builder
2.创建网络定义: builder --> network
3.配置参数: builder --> config
4.生成engine: builder --> engine(network, config)
5.序列化保存: engine --> serialize
6.释放资源: delete
*/

#include<iostream>
#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));                              // 写入 权重size(大小: 代表数据的多少), 传入地址
    outfile.write((char *)(data), size * sizeof(float));                      // 写入 权重data, 传入地址
    outfile.close();
   
}

// 读取权重, vector 进行读取
std :: vector<float> loadWeights(const std :: string &filename)
{
    // 读取文件
    std :: ifstream infile(filename, std :: ios :: binary);
    assert(infile.is_open() && "load weights failed");                          // assert 断言, 如果条件不满足就会报错

    int size;
    infile.read((char *)(&size), sizeof(int));                                  // 读取 权重size(大小: 代表数据的多少), 传入地址

    std :: vector<float> data(size);                                            // 权重data, 传入大小
    infile.read((char *)(data.data()), size * sizeof(float));                   // 读取 权重data, 传入地址
    infile.close();

    return data;                                                                // 返回形式是: vector 

}


int main()
{
    // ====================== 1. 创建 builder ======================
    TRTLogger logger;
    nvinfer1 :: IBuilder* builder = nvinfer1 :: createInferBuilder(logger);


    // ====================== 2. 创建网络定义 ======================

    // 调用createNetworkV2()创建网络结构, 1: 代表显性 batch
    nvinfer1 :: INetworkDefinition *network = builder -> createNetworkV2(1);

    // 定义网络结构
    // mlp 多层感知机: input(1, 3, 1, 1) --> fc1 --> sigmoid --> output(2)

    // =========== 创建一个 input tensor ===========
    const int input_size = 3;
    nvinfer1 :: ITensor *input = network -> addInput("data", nvinfer1 :: DataType :: kFLOAT, nvinfer1 :: Dims4(1, input_size, 1, 1));

    // =========== 创建全连接层 fc1 ===========

    // weight 和 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};

    // 将 weight 和 bias 转为 nvinfer 模型, 参数分别是: data_type, data(指针获取数据传过来是地址), size(大小, 个数)
    // nvinfer1 :: Weights fc1_weights {nvinfer1 :: DataType :: kFLOAT, fc1_weight_data, input_size * 2};
    // nvinfer1 :: Weights fc1_bias {nvinfer1 :: DataType :: kFLOAT, fc1_bias_data, 2};

    // ====== 将 weight 和 bias 保存权重 ======
    // 将权重保存到文件中, 演示从别的来源加载权重, 
    saveWeights("./model/fc1.wts", fc1_weight_data, 6);
    saveWeights("./model/fc1.bias", fc1_bias_data, 2);

    // 读取对应的权重, auto 进行读取, loadWeights()函数返回的是 vec
    auto fc1_weights_vec = loadWeights("./model/fc1.wts");
    auto fc1_bias_vec = loadWeights("./model/fc1.bias");

    // ====== 从保存权重获取weight和bias ======
    // 将 weight 和 bias 转为 nvinfer 模型, 参数分别是: data_type, data(vec获取数据是 data), size(大小, 个数)
    nvinfer1 :: Weights fc1_weight {nvinfer1 :: DataType :: kFLOAT, fc1_weights_vec.data(), fc1_weights_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 ===========
    // 添加 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. 配置参数: builer ---> config ======================

    // 添加配置参数, 告诉 tensorRT 应该如何优化网络
    nvinfer1 :: IBuilderConfig *config = builder -> createBuilderConfig();

    // 设置最大工作空间, 单位是字节
    config -> setMaxWorkspaceSize(1 << 28);                                  // 1 << 28: 256MiB


    // ====================== 4. 生成engine：builder --> network --> config ======================
    nvinfer1 :: ICudaEngine *engine = builder -> buildEngineWithConfig(*network, *config);

    // 如果没有生成 engine, 就简单的报错一下
    if(!engine)
    {
        std :: cout << "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;

}

runtime.cu

// TensorRT运行时的最高层级接口是Runtime
// 执行推理的部分

/*
使用 .cu 是希望使用 CUDA 的编译器 NVCC, 会自动连接 cuda 库

TensorRT runtime 的推理过程
1. 创建一个 runtime 对象
2. 反序列化生成 engine: runtime --> engine
3. 创建一个执行上下文 (ExecutionContext调用enqueueV2()来运行Inference) ExecutionContext: engine --> context


    4. 填充数据
    5. 执行推理: context --> enqueueV2

6. 释放资源: delete

*/


#include<iostream>
#include<vector>
#include<cassert>
#include<fstream>

#include<NvInfer.h>
#include<cuda_runtime.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;


// 加载模型: mlp.engine
std :: vector<unsigned char> loadEngineModel(const std :: string &filename)
{
    std :: ifstream infile(filename, std :: ios :: binary);                       // 以二进制方式读取文件
    
    // 断言
    assert(infile.is_open() && "load weights failed");
    
    file.seekg(0, std :: ios :: binary);                            // 定位到文件末尾
    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);                         //
    
    file.close();

    return data;
   
}




int main()
{
    // ========================== 1. 创建一个 runtime 对象 ==========================
    // TRTLogger 实例化
    TRTLogger logger;
    nvinfer1 :: IRuntime *runtime = nvinfer1 :: createInferRuntime(logger);


    // ========================== 2. 反序列化生成 engine: runtime --> engine ==========================
    // 读取文件
    auto engineModel = loadEngineModel("./model/mlp.engine");

    // 调用 runtime 的反序列方法, 生成 engine, 参数: 模型数据地址, 模型大小, pluginFactory
    nvinfer1 :: ICudaEngine *engine = runtime -> deserializeCudaEngine(engineModel.data(), engineModel.size(), nullprt);

    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);

    // cpu: host  GPU: device
    // 数据流转: 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[3] {0, 0,};                            // host 输出数据, 初始化0, 0
    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;

    // 释放资源
    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;

}