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: 代表显性 batchnvinfer1 :: 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()函数返回的是 vecauto 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, biasnvinfer1 :: 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_sizebuilder -> 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 --> context4. 填充数据5. 执行推理: context --> enqueueV26. 释放资源: 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, 大小为 sizefile.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, 参数: 模型数据地址, 模型大小, pluginFactorynvinfer1 :: 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, 0int 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;// 释放资源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;}