5_TensorRT实战-基本框架

上篇博文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 mEngine = nullptr; //用于运行网络TensorRT引擎
	
	    //!
	    //! 该函数为GoogleNet解析一个Caffe模型，并创建一个TensorRT网络
	    //!
	    void constructNetwork(SampleUniquePtr& builder, SampleUniquePtr& network, SampleUniquePtr& 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::createInferBuilder(gLogger));

    if (!builder)
        return false;
	
	//通过builder创建网络定义
    auto network = SampleUniquePtr(builder->createNetwork());
    if (!network)
        return false;
	
	//创建用于解析caffe网络模型的parser
    auto parser = SampleUniquePtr(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(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(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网络时可能借鉴这一思想。