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<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網絡時可能借鑑這一思想。