工作流程
代码:
// 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));
//weight需要count(个数),type(数据类型)和value(数值)确定
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)));
//1<<28 1左移28位相当于*2的28次方即256mb
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;
}
整体流程为:
1. 定义builder,config,network:
builder:
nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
config:
nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
network:
ncinfer1::INetworkDefinition* network = builder->createNetworkV2(1);
2. 输入层,全连接层,全连接参数(w,b)以及激活层
input:addInput(名称,数据类型,形状(nchw))
nvinfer1::ITensor* input =
network->addInput("image",nvinfer1::DataType::KFLOAT,nvinfer1::Dims4(N,C,H,W);
fullyconection
auto layer1 = network->addFullyConnected(*input,channel_out,layer1_weight,layer1_bias);
activation:
auto prob = network.addActivation(*(layer1->getOutput(0)),nvinfer1::ActivationType::KSIGMOID);
3.标记最终输出
network->markOutput(*(prob->getOutput(0)));
4.内存配置和batchsize设置
config->setMaxWorkspaceSize(1<<28);//256mb
builder->setMaxBatchSize(1);//bs=1
5. 构建engine
nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network,config);
6. 转trt文件
//序列化
nvinfer1::IHostMemory* model_data = engine->serialize();
//创建文件
FILE* f = fopen("engine.trtmodel","wb");
//写入
fwrite(model_data->data(),1,model_data->size(),f);
fclose();
重点:
1. 创建network使用createtworkV2,并指定batchs_size显式为1,后续推理用enqueueV2
2. builder、config等记得释放,否则会有内存泄漏,用ptr->destroy()
3. markOutput表示是该模型的输出节点,mark几次就有几个输出,addInput几次就有几个输入
4. workspace供相同layer内存复用
5. 保存的模型只能适配编译的trt版本、设备,如果跨平台,跨版本都需要重新编译保证可用和最优配置。