【TensorRT】使用C++ API直接创建tensorRT 模型

Linnea ·
更新时间:2024-09-21
· 612 次阅读

使用TensorRT C++ API创建的模型只有前向推理能力,没有反向传播能力。也就是说不能用于训练,模型的权值参数要人为给定。但可以针对设定网络模型(自己使用API创建网络模型)或给定模型(使用NvCaffeParser或NvUffPaser导入其他深度学习框架训练好的模型)做一系列优化,以加快推理速度(inference)。

使用C++ API函数部署网络主要分为四个步骤:

创建网络; 为网络添加输入; 添加各种各样的层; 设定网络输出。

这里只列出一些关键代码,详细的代码可以参考TensorRT官方的例子。

一、创建网络

首先创建一个TensorRT的network,这个network目前没有什么实质的内容。

nvinfer1::IBuilder* = nvinfer1::createInferBuilder(gLogger.getTRTLogger()); INetworkDefinition* network = builder->createNetwork(); 二、为网络添加输入

所有的网络都必须明确输入是哪个blob,TensorRT才能知道数据的入口。

data = network->addInput(INPUT_BLOB_NAME, dt, DimsCHW{ 1, INPUT_H, INPUT_W});

其中,INPUT_BLOB_NAME 是为输入 blob起的名字。

dt是指数据类型,我使用的TensorRT版本是6.0.1.5,关于数据类型的定义位于include/NvInferRuntimeCommon.h文件内。

enum class DataType : int { kFLOAT = 0, //!< FP32 format. kHALF = 1, //!< FP16 format. kINT8 = 2, //!< quantized INT8 format. kINT32 = 3 //!< INT32 format. };

DimsCHW{ 1, INPUT_H, INPUT_W},其中batch为1(省略),channel 为1,输入height 和width分别为 INPUT_H, INPUT_W的blob。

三、为网络添加层

可以添加不包含训练参数的层,比如Relu层,Pooling层等,和包含训练参数的层,比如卷积层,全连接层等。

3.1 添加不包含训练参数的层

比如这里添加一个对输入数据做归一化的scale层。

const float scaleParam = 0.0125f; const Weights power{DataType::kFLOAT, nullptr, 0}; const Weights shift{DataType::kFLOAT, nullptr, 0}; const Weights scale{DataType::kFLOAT, &scaleParam, 1}; IScaleLayer* scale_1 = network->addScale(*data, ScaleMode::kUNIFORM, shift, scale, power); 3.2 添加包含训练参数的层

一般来说包含训练参数的层的参数是从已经训练好的模型文件中得来的,所以要先加载模型文件。

比如这里添加一个卷积层。

std::map weightMap = loadWeights(locateFile("mnistapi.wts")); IConvolutionLayer* conv1 = network->addConvolution(*scale_1->getOutput(0), 20, DimsHW{5, 5}, weightMap["conv1filter"], weightMap["conv1bias"]); conv1->setStride(DimsHW{1, 1});

其中*scale_1->getOutput(0)是获取上一层 scale层的输出。

20代表的是卷积核的个数。

DimsHW{5, 5}是卷积核大小5x5。

weightMap["conv1filter"], weightMap["conv1bias"]是从加载的模型文件得到的权值系数矩阵。

四、设定网络输出

网络必须指定哪一个blob是输出的。

auto prob = network->addSoftMax(*ip2->getOutput(0)); prob->getOutput(0)->setName(OUTPUT_BLOB_NAME); network->markOutput(*prob->getOutput(0));

这里最后为网络添加了一个softmax层,并将其命名为OUTPUT_BLOB_NAME,最后将这个层指定为输出层。


作者:heiheiya



c+ tensorrt 模型 C++ api

需要 登录 后方可回复, 如果你还没有账号请 注册新账号