网站导航
> 文章列表 > 第一个TensorRT程序，写一个多层感知机

第一个TensorRT程序，写一个多层感知机

网友: 文章列表 2024-03-21 16:47:16

第一个TensorRT程序，写一个多层感知机

TensorRT build engine的流程

  1. 创建builder:
  2. 创建网络定义 builder --> network
  3. 配置参数: builder --> config
  4. 生成engine: builder --> engine(network, config)
  5. 序列化保存: engine --> serialize
  6. 释放资源: delete

1. 创建builder

#include <iostream>
#include <NvInfer.h>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;int main() {// 1. 创建builderTRTLogger logger; // logger是必要的，用来捕捉warning和info等nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);return 0;
}

nvinfer1下的数据类型，都是nvinfer1::Data ptr, 后面通过指针操作

Severity 是一个枚举类型，用于指定日志消息的严重程度。在 TensorRT 中，Severity 枚举类型定义了四个常量，分别为：

kINTERNAL_ERROR：内部错误，表示程序出现了无法处理的错误或异常。

kERROR：错误，表示程序执行过程中出现了错误，但程序可以继续运行。

kWARNING：警告，表示程序执行过程中出现了一些可能会导致问题的情况。

kINFO：信息，表示程序执行过程中的一些有用的信息。

在使用 TensorRT 进行模型推理时，我们可以使用 ILogger 接口来定义日志输出，例如将日志输出到控制台或文件中。在实现 ILogger 接口时，需要实现 log() 函数来处理日志消息，该函数的参数包括一个 Severity 类型的参数和一个 const char* 类型的参数，分别表示日志消息的严重程度和内容。

在示例程序中，我们定义了一个名为 TRTLogger 的类，它是 ILogger 接口的实现类。在 TRTLogger 类中，我们通过重载 log() 函数来过滤掉 Severity 为 kINFO 的日志消息，并将其他日志消息输出到控制台中。这样可以避免输出过多无用的信息，从而更好地了解程序的运行情况。

1.1 复习抽象基类和静态函数

nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);

首先这里的nvinfer1::IBuilder 是一个抽象基类, createInferBuilder()是一个静态函数, 抽象基类需要new开辟空间

#include <iostream>class Shape {
public:virtual void draw() = 0;  // 纯虚函数static Shape* createShape(const std::string& type);
};class Rectangle : public Shape {
public:void draw() override {std::cout << "Drawing a rectangle." << std::endl;}
};class Circle : public Shape {
public:void draw() override {std::cout << "Drawing a circle." << std::endl;}
};Shape* Shape::createShape(const std::string& type) {if (type == "Rectangle") {return new Rectangle();} else if (type == "Circle") {return new Circle();} else {return nullptr;}
}int main() {Shape* s1 = new Rectangle();Shape* s2 = Shape::createShape("Circle");s1->draw();s2->draw();delete s1;delete s2;return 0;
}

在这段代码中，Shape 是一个抽象基类，它包含一个纯虚函数 draw()，这意味着该函数没有实现，子类必须覆盖该函数并提供自己的实现。这样做的目的是将 Shape 定义为一个通用概念，但由于它本身是不完整的，因此不能实例化。通过将 Shape 设计为抽象基类，它的子类必须实现 draw()，从而保证了 Shape 的通用性和可扩展性。

Shape 的子类 Rectangle 和 Circle 继承了 Shape 类，并实现了 draw() 函数。Rectangle 类的 draw() 函数打印“Drawing a rectangle.”，而 Circle 类的 draw() 函数打印“Drawing a circle.”。

Shape::createShape() 是一个静态函数，它接受一个 std::string 类型的参数 type，并返回一个 Shape 类型的指针。在此示例中，它检查 type 的值，如果为 “Rectangle”，则创建一个 Rectangle 对象，如果为 “Circle”，则创建一个 Circle 对象。如果 type 的值无法识别，则返回 nullptr。这个函数的目的是封装对象的创建，并提供一种灵活的方式来创建不同的对象，而无需暴露对象创建的具体细节。在这个例子中，createShape() 方法是静态的，这意味着可以直接通过类名来调用该方法，而无需先创建类的对象。

2. 定义网络

  1. build -> network
  • 这里用显性batch, 显性 batch 是指在网络定义时指定 batch size 的大小。这是一种固定的 batch 大小，不会随输入数据而变化，而且可以通过 TensorRT 进行优化。
  • 显性 batch 是指在网络定义时指定 batch size 的大小。这是一种固定的 batch 大小，不会随输入数据而变化，而且可以通过 TensorRT 进行优化。
  1. net->addInput()
  • KFLOAT, KHALF, KINT8, KINT32 是 TensorRT 中用于表示张量数据类型的枚举类型
  • Dim4{1, 3, 1, 1} batch, channel, height_Size, width_size
  1. network->addFullyConnected()
  • input, weight, bias是指针，
  • 这三个参数都需要被转成API需要的数据类型
  • network->addInput() nvinfer1::Weights

在 TensorRT 中，每个层都会产生一个或多个输出张量。对于单个输出张量的层，例如全连接层和卷积层，可以通过 getOutput(0) 方法来获取输出张量。对于产生多个输出张量的层，例如多输出卷积层和循环神经网络层，可以通过指定输出张量的索引来获取对应的输出张量，例如 getOutput(0)、getOutput(1) 等。在这个例子中，我们可以通过 fc1->getOutput(0) 来获取全连接层的输出张量，通过 sigmoid->getOutput(0) 来获取 Sigmoid 激活层的输出张量。这样做是为了方便在构建网络的过程中访问中间结果。

如果我们有多个输出，而只标记其中的一个，那么其他的输出将可能不被优化，从而导致不必要的计算和额外的延迟。因此，在多输出情况下，我们需要显式地标记所有的输出。

 // 2. 创建网络定义, builder --> network, 显性的batchnvinfer1::INetworkDefinition *network = builder->createNetworkV2(1);/*定义网络结构: 多层感知机Input: (1, 3, 1, 1) -> fc -> sigmoid() -> output(2)network->addInput(name, DataType, Dims)nvinfer1::Weights obj{DataType, ptr, size}*/const int input_size = 3;nvinfer1::ITensor *input = network->addInput("data", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, input_size, 1, 1});// weight and bias, 都是在堆上开辟内存返回指针, 满足API需求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};// 把堆上开辟的参数转换成nvinfer1::Weights类型， nvinfer1::Weights obj{DataType, ptr, size}nvinfer1::Weights fc1_weight{nvinfer1::DataType::kFLOAT, fc1_weight_data, input_size * 2};nvinfer1::Weights fc1_bias{nvinfer1::DataType::kFLOAT, fc1_bias_data, 2};const int output_size = 2;// nvinfer1::IFullyConnectedLayer 添加全连接层nvinfer1::IFullyConnectedLayer *fc1 = network->addFullyConnected(*input, output_size, fc1_weight, fc1_bias);// nvinfer1::IActivationLayer 添加激活层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

// ======================3. 配置参数=============================
// 添加配置参数, 告诉TensorRT如何优化网络
nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
// 设置最大workspace, 单位是字节
config->setMaxWorkspaceSize(1<<28);  // 256MiB = 2^28 / 1024 / 1024

通过添加配置参数，可以告诉TensorRT如何去优化网络，以达到更好的推理性能。比如，可以设置最大的workspace大小、允许的最大批量大小、推理精度、允许的算子合并等等，这些参数都可以影响TensorRT对网络的优化方式和效果。因此，合理的配置参数可以在保证模型精度的前提下，最大化TensorRT的推理速度和性能。

workspace是用来在优化网络时分配内存使用的。在实际应用中，需要根据网络规模和硬件资源设置合适的workspace大小。如果设置过小，可能会导致内存不足而无法优化网络；如果设置过大，可能会浪费内存资源。

在多次优化同一网络时，如果网络规模没有变化，workspace大小也可以不必每次都设置。但如果网络结构有较大变化，最好重新设置一下workspace大小以适应新的网络结构。

4. 创建engine和序列化engine

// ======================4. 创建engine: builder ->network -> config===================
nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
if (!engine)
{std::cerr << "创建Engine失败! " << std::endl;return -1;
}// ===================5. 序列化engine=======================
// 保存到磁盘上
nvinfer1::IHostMemory *serialized_engine = engine->serialize();
// 存入文件
std::ofstream outfile("engine", std::ios::binary);
assert(outfile.is_open() && "Failed to open file for writting");
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;

创建engine需要传入*network, *config。这两个是nvinfer1下的数据类型, 他们的指针是network, config。然后定义一个二进制文件engine, 写入两个参数, engine的data和engine的size()再释放

5. 完整的构建engine的main.cpp文件

#include <iostream>
#include <fstream>
#include <NvInfer.h>
#include <cassert>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;int main() {// =======================1. 创建builder==============================TRTLogger logger; // logger是必要的，用来捕捉warning和info等nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(logger);//  =======================2. 创建网络定义==============================// builder --> network, 显性的batchnvinfer1::INetworkDefinition *network = builder->createNetworkV2(1);/*定义网络结构: 多层感知机Input: (1, 3, 1, 1) -> fc -> sigmoid() -> output(2)network->addInput(name, DataType, Dims)nvinfer1::Weights obj{DataType, ptr, size}*/const int input_size = 3;nvinfer1::ITensor *input = network->addInput("data", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, input_size, 1, 1});// weight and bias, 都是在堆上开辟内存返回指针, 满足API需求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};// 把堆上开辟的参数转换成nvinfer1::Weights类型， nvinfer1::Weights obj{DataType, ptr, size}nvinfer1::Weights fc1_weight{nvinfer1::DataType::kFLOAT, fc1_weight_data, input_size * 2};nvinfer1::Weights fc1_bias{nvinfer1::DataType::kFLOAT, fc1_bias_data, 2};const int output_size = 2;// nvinfer1::IFullyConnectedLayer 添加全连接层nvinfer1::IFullyConnectedLayer *fc1 = network->addFullyConnected(*input, output_size, fc1_weight, fc1_bias);// nvinfer1::IActivationLayer 添加激活层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. 配置参数=============================// 添加配置参数, 告诉TensorRT如何优化网络nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();// 设置最大workspace, 单位是字节config->setMaxWorkspaceSize(1<<28);  // 256MiB = 2^28 / 1024 / 1024// ======================4. 创建engine: builder ->network -> config===================nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);if (!engine){std::cerr << "创建Engine失败! " << std::endl;return -1;}// ===================5. 序列化engine=======================// 保存到磁盘上nvinfer1::IHostMemory *serialized_engine = engine->serialize();// 存入文件std::ofstream outfile("engine", std::ios::binary);assert(outfile.is_open() && "Failed to open file for writting");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;
}

6. runfile.cu文件

  1. 创建runfile
  2. 反序列化engine
  3. 创建context
  4. 导入推理的数据, 放到GPU上面

在这里context 用于执行一次推理。首先，我们为输入数据和输出数据分配了 GPU 内存。然后，我们将主机上的输入数据传输到 GPU 上的内存中。接着，我们定义了一个输入输出数组 bindings，告诉 context 输入和输出数据的位置。最后，我们调用了 context 的 enqueueV2 方法来执行推理操作，并等待推理操作完成。推理完成后，我们将输出数据从 GPU 上的内存传输到主机上，最后释放资源，包括 GPU 内存、context 对象、engine 对象、runtime 对象等。

#include <iostream>
#include <vector>
#include <fstream>
#include <cassert>
#include "cuda_runtime.h"
#include "NvInfer.h"// logger用来管控打印日志级别
class TRTLogger : public nvinfer1::ILogger
{void log(Severity severity, const char *msg) noexcept override{// 屏蔽INFO级别的日志if (severity != Severity::kINFO)std::cout << msg << std::endl;}
};// 这里使用无符号整数, size_t文件大小不会是负数, unsigned char是避免溢出时出现负数
std::vector<unsigned char> loadEngineModel(const std::string &fileName)
{std::ifstream file(fileName, std::ios::binary);assert(file.is_open() && "load engine model Failed");// 创建data(size)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. =============1. 创建一个runtime对象===========================TRTLogger logger;nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);// ==============2. 反序列化生成engine============================// 读取文件auto engineModel = loadEngineModel("engine");// 调用runtime的反序列化方法，生成engine，参数分别是：模型数据地址，模型大小，pluginFactorynvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(engineModel.data(), engineModel.size(), nullptr);if (!engine){std::cout << "反序列化Engine失败" << std::endl;}// ==================3. 创建一个执行的上下文==========================// 创建contextnvinfer1::IExecutionContext *context = engine->createExecutionContext();cudaStream_t stream = nullptr;cudaStreamCreate(&stream);// 输入数据float *host_input_data = new float[3]{2, 4, 8}; // host 输入数据int input_data_size = 3 * sizeof(float);float *device_input_data = nullptr;// 输出数据float *host_output_data = new float[2]{0, 0};int output_data_size = 2 * sizeof(float);float *device_output_data = nullptr;// GPU开辟内存cudaMalloc((void **)&device_input_data, input_data_size);cudaMalloc((void **)&device_output_data, output_data_size);// 主机数据分配到device上cudaMemcpyAsync(device_input_data, host_input_data, input_data_size, cudaMemcpyHostToDevice, stream);// bindings告诉context输入输出数据的位置float *bingdings[] = {device_input_data, device_output_data};// ============5. 执行推理=================bool sucess = context->enqueueV2((void **)bingdings, 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;
}

7. 总结

在主程序(main.cpp)中，我们定义了网络结构和权重，并使用Builder和Config对象来设置和优化网络。最终，我们使用Builder对象的buildEngineWithConfig()函数将网络和配置参数组合在一起，并将生成的engine序列化到磁盘中，以便在推理时使用。

在推理程序(runfile.cu)中，我们使用IRuntime对象的deserializeCudaEngine()函数从磁盘中反序列化engine。然后，我们使用Context对象将输入和输出绑定到相应的GPU缓冲区上，并使用enqueueV2()函数启动异步推理。在推理完成后，我们使用CUDA流和异步内存拷贝函数将结果从GPU缓冲区复制到主机内存中，最后输出结果。

特殊符号

网络标签：

相关问题