目录
前言
杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。
本次课程学习 tensorRT 基础-第一个 trt 程序，实现模型编译的过程
课程大纲可看下面的思维导图
1. hello案例
学习使用 TensorRT-CPP 的 API 构建模型，并进行编译的流程
案例代码如下：
// tensorRT include
#include <NvInfer.h>
#include <NvInferRuntime.h>
// cuda include
#include <cuda_runtime.h>
// system include
#include <stdio.h>
class TRTLogger : public nvinfer1::ILogger{
public:
virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override{
if(severity <= Severity::kVERBOSE){
printf("%d: %s\n", severity, msg);
}
}
};
nvinfer1::Weights make_weights(float* ptr, int n){
nvinfer1::Weights w;
w.count = n;
w.type = nvinfer1::DataType::kFLOAT;
w.values = ptr;
return w;
}
int main(){
// 本代码主要实现一个最简单的神经网络 figure/simple_fully_connected_net.png
TRTLogger logger; // logger是必要的，用来捕捉warning和info等
// ----------------------------- 1. 定义 builder, config 和network -----------------------------
// 这是基本需要的组件
//形象的理解是你需要一个builder去build这个网络，网络自身有结构，这个结构可以有不同的配置
nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
// 创建一个构建配置，指定TensorRT应该如何优化模型，tensorRT生成的模型只能在特定配置下运行
nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
// 创建网络定义，其中createNetworkV2(1)表示采用显性batch size，新版tensorRT(>=7.0)时，不建议采用0非显性batch size
// 因此贯穿以后，请都采用createNetworkV2(1)而非createNetworkV2(0)或者createNetwork
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);
// 构建一个模型
/*
Network definition:
image
|
linear (fully connected) input = 3, output = 2, bias = True w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5]], b=[0.3, 0.8]
|
sigmoid
|
prob
*/
// ----------------------------- 2. 输入，模型结构和输出的基本信息 -----------------------------
const int num_input = 3; // in_channel
const int num_output = 2; // out_channel
float layer1_weight_values[] = {
1.0, 2.0, 0.5, 0.1, 0.2, 0.5}; // 前3个给w1的rgb，后3个给w2的rgb
float layer1_bias_values[] = {
0.3, 0.8};
//输入指定数据的名称、数据类型和完整维度，将输入层添加到网络
nvinfer1::ITensor* input = network->addInput("image", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4(1, num_input, 1, 1));
nvinfer1::Weights layer1_weight = make_weights(layer1_weight_values, 6);
nvinfer1::Weights layer1_bias = make_weights(layer1_bias_values, 2);
//添加全连接层
auto layer1 = network->addFullyConnected(*input, num_output, layer1_weight, layer1_bias); // 注意对input进行了解引用
//添加激活层
auto prob = network->addActivation(*layer1->getOutput(0), nvinfer1::ActivationType::kSIGMOID); // 注意更严谨的写法是*(layer1->getOutput(0)) 即对getOutput返回的指针进行解引用
// 将我们需要的prob标记为输出
network->markOutput(*prob->getOutput(0));
printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f); // 256Mib
config->setMaxWorkspaceSize(1 << 28);
builder->setMaxBatchSize(1); // 推理时 batchSize = 1
// ----------------------------- 3. 生成engine模型文件 -----------------------------
//TensorRT 7.1.0版本已弃用buildCudaEngine方法，统一使用buildEngineWithConfig方法
nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
if(engine == nullptr){
printf("Build engine failed.\n");
return -1;
}
// ----------------------------- 4. 序列化模型文件并存储 -----------------------------
// 将模型序列化，并储存为文件
nvinfer1::IHostMemory* model_data = engine->serialize();
FILE* f = fopen("engine.trtmodel", "wb");
fwrite(model_data->data(), 1, model_data->size(), f);
fclose(f);
// 卸载顺序按照构建顺序倒序
model_data->destroy();
engine->destroy();
network->destroy();
config->destroy();
builder->destroy();
printf("Done.\n");
return 0;
}
运行效果如下：
上述示例代码演示了使用 tensorRT 构建一个简单神经网络模型的过程。通过定义模型结构、设置输入和输出信息，并使用 TensorRT 的 Builder 对象和网络定义对象构建 Engine 模型。然后，将生成的模型序列化并保存为文件。示例代码中还展示了如何配置 TensorRT 的优化参数和限制条件。通过这个示例代码，可以了解 TensorRT 构建模型的基本工作流程和相关对象的使用方法。
首先引入了必要的头文件和库，代码中引入了 TensorRT 和 CUDA 的相关头文件和库。然后定义了一个日志类，继承自 nvinfer1::ILogger，用于捕获 TensorRT 的警告和信息。
定义输入、模型结构和输出的基本信息：指定输入张量的名称、数据类型和维度，添加全连接层和激活层，并将输出标记为网络的输出。接下来配置 TensorRT，通过设置最大工作空间大小和最大批处理大小来配置，随后构建 Engine 模型，使用 buildEngineWithConfig 方法构建 Engine 模型。序列化并保存模型，将模型序列化为二进制数据，并将其保存为文件。最后释放资源，按照创建顺序的倒序释放 TensorRT 的相关资源。
TensorRT 构建模型的大致工作流程可以分为四个部分：
- 1.定义 builder，config 和 network
- 2.构建网络所需输入，模型结构和输出的基本信息
- 3.生成 engine 模型文件
- 4.序列化模型文件并存储
关于该示例代码的重点提炼：
- 必须使用 createNetworkV2，并指定为 1（表示显性 batch）。createNetwork 已经废弃，非显性 batch 官方不推荐。这个方式直接影响推理时 enqueue 还是 enqueueV2
- builder、config 等指针记得释放，否则会有内存泄露，使用 ptr->destroy() 释放
- markOutput 表示该模型的输出节点，mark 几次就有几个输出，addInput 几次就有几个输入，这与推理时相呼应
- workspaceSize 是工作空间大小，某些 layer 需要使用额外存储时，不会自己分配空间，而是为了内存复用，直接找 tensorRT 要 workspace 空间。指的这个意思
- 一定要记住，保存的模型只能适配编译时的 trt 版本、编译时指定的设备。也只能保证在这种配置下是最优的。如果用 trt 跨不同设备执行，有时可以运行，但不是最优的，也不推荐
2. 补充知识
关于第一个 trt 程序的相关知识点：(from 杜老师)
- main.cpp 构建了一个最简单全连接网络
- tensorrt 的工作流程如下图：
- 首先定义网络
- 优化 builder 参数
- 通过 builder 生成 engine 用于模型保存、推理等
- engine 可以通过序列化和反序列化转化模型数据类型（转化为二进制 byte 文件，加快传输速率），再一步推动模型由输入张量到输出张量的推理
- code struct
- 1.定义 builder, config 和network，其中 builder 表示所创建的构建器，config 表示创建的构建配置（指定 TensorRT 应该如何优化模型），network 为创建的网络定义。
- 2.输入，模型结构和输出的基本信息（如下图所示）
- 3.生成 engine 模型文件
- 4.序列化模型文件并存储
总结
本次课程学习了使用 tensorRT 的 C++ 接口来搭建一个简单的神经网络结构，整体流程可分为：builder、config、network 定义；输入、模型结构和输出信息；engine 模型文件生成；序列化模型文件并存储四个部分。
一些细节需要大家自行看代码进行分析，比如 builder、config 等指针记得释放，在 tensorRT 构建网络前需要定义日志 Logger 类用于捕获 tensorRT 的信息等等