上篇博文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网络时可能借鉴这一思想。
