上篇博文4_TensorRT概况主要讲了Nvida TensorRT的编程API，本篇主要通过一个简单、完整的例子来讲解如何将一个Caffe模型(GoogleNet模型)通过TensorRT进行推理加速。
系统环境
本示例运行的系统环境如下：
-
硬件环境：Jetson TX2
-
软件环境:
- JetPack:V4.2
- CUDA:CUDA ToolKit for L4T V10.0
- cuDNN:
- cuDNN on Target 7.3
- TensorRT On Target 5.0
- Computer Vison:
- OpenCV on Target 3.3.1
- VisionWorks on target 1.6
- MultiMedia API: 32.1
TensorRT基本框架
SampleGoogleNet类实现了基于GoogleNet模型的TensorRT网络构建、engine创建、推理等接口。
class SampleGoogleNet
{
public:
SampleGoogleNet(const samplesCommon::CaffeSampleParams& params)
: mParams(params)
{
}
//!
//! 创建TensorRT网络
//!
bool build();
//!
//! 运行TensorRT推理引擎
//!
bool infer();
//!
//! 清理运行时创建的状态、资源
//!
bool teardown();
samplesCommon::CaffeSampleParams mParams;
private:
std::shared_ptr<nvinfer1::ICudaEngine> mEngine = nullptr; //用于运行网络TensorRT引擎
//!
//! 该函数为GoogleNet解析一个Caffe模型，并创建一个TensorRT网络
//!
void constructNetwork(SampleUniquePtr<nvinfer1::IBuilder>& builder, SampleUniquePtr<nvinfer1::INetworkDefinition>& network, SampleUniquePtr<nvcaffeparser1::ICaffeParser>& parser);
};
配置参数
构建TensorRT时，需要几个比较重要的参数，这些参数一般在TensorRT应用启动时由命令行传入或者使用默认的配置参数。大部分参数都是在构建TensorRT网络时需要的配置参数，先分别列出如下：
- batchSize:批量输入的数量
- dalCore:是否使用DLA(Deep Learning Accelerate
- dataDirs:网络模型数据存放的位置
- inputTensorNames:用作输入的Tensor的数量
- outputTensorNames:用作输出的Tensor的数量
- 下面这两个参数用于基于Caffe的神经网络配置：
- prototxtFileName:网络原型配置文件
- weightsFileName:网络权重文件
构建（网络、推理引擎）
SampleGoogleNet::build(),该函数通过解析caffe模型创建GoogleNet网络，并构建用于运行GoogleNet (mEngine)的引擎。
//创建用于推理的Builder
auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(gLogger));
if (!builder)
return false;
//通过builder创建网络定义
auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetwork());
if (!network)
return false;
//创建用于解析caffe网络模型的parser
auto parser = SampleUniquePtr<nvcaffeparser1::ICaffeParser>(nvcaffeparser1::createCaffeParser());
if (!parser)
return false;
//通过builder、network、parser、配置参数构建网络定义
constructNetwork(builder, network, parser);
constructNetwork函数定义如下：
{
const nvcaffeparser1::IBlobNameToTensor* blobNameToTensor = parser->parse(
locateFile(mParams.prototxtFileName, mParams.dataDirs).c_str(),//加载网络原型配置文件
locateFile(mParams.weightsFileName, mParams.dataDirs).c_str(),//加载网络训练权重文件
*network,//网络定义
nvinfer1::DataType::kFLOAT);//权重和张量的精度类型为FP32 format
//遍历outputTensorNames，通过blobNameToTensor->find函数转换为对应的Tensor，最后通过markOutput将该Tensor标记为网络的输出量。
for (auto& s : mParams.outputTensorNames)
network->markOutput(*blobNameToTensor->find(s.c_str()));
//根据batchSize设置最大的batchsize。
builder->setMaxBatchSize(mParams.batchSize);
//设置最大的工作空间大小。
builder->setMaxWorkspaceSize(16_MB);
//根据dlaCore决定是否启用DLA功能。
samplesCommon::enableDLA(builder.get(), mParams.dlaCore);
}
//根据构建好的网络定义创建Cuda推理引擎。
mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(builder->buildCudaEngine(*network), samplesCommon::InferDeleter());
if (!mEngine)
return false;
推理
SampleGoogleNet::infer()，这个函数是示例的主要执行函数。它分配缓冲区、设置输入并执行引擎。
//创建RAII缓冲区(BufferManager类处理主机和设备(GPU)缓冲区分配和释放)管理结构。
//BufferManager这个RAII类处理主机和设备缓冲区的分配和释放，主机和设备缓冲区之间的memcpy来辅助推理，调试转储来验证推
//理。BufferManager类用于简化缓冲区管理以及缓冲区与引擎之间的任何交互
samplesCommon::BufferManager buffers(mEngine, mParams.batchSize);
//创建推理引擎运行上下文
auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
if (!context)
return false;
//获取主机缓冲区并将主机输入缓冲区设置为所有零
for (auto& input : mParams.inputTensorNames)
memset(buffers.getHostBuffer(input), 0, buffers.size(input));
//将数据通过memory从主机输入缓冲区拷贝到设备输入缓冲区
buffers.copyInputToDevice();
//执行推理
bool status = context->execute(mParams.batchSize, buffers.getDeviceBindings().data());
if (!status)
return false;
//推理完成之后，将数据通过memcopy从设备输出缓冲区拷贝到主机输出缓冲区
buffers.copyOutputToHost();
资源清理
nvcaffeparser1::shutdownProtobufLibrary();资源清理主要涉及到parser所使用的protobuf的清理。
总结
本文通过一个十分简单的示例，讲解了如何将一个网络模型部署到TensorRT上的编码过程。需要注意的是，文中所使用的网络模型为Caffe，使用到的parser也为ICaffeParser，TensorRT同时还支持ONX、UFF格式的parser，后续会总结如何通过这两类parser导入其他不同的网络模型，例如,tensorflow等。在执行推理时，需要涉及到数据在GPU缓存与CPU缓存之间的拷贝过程，该过程比较繁琐，文中使用BufferMannager很好的封装了这些过程，后续再开发TensorRT网络时可能借鉴这一思想。
