Tensorrt使用入门（环境搭建与C++ API调用）

最近在学习tensorrt，写篇博客记录一下，本博客只针对入门使用者。

TensorRT介绍

NVIDIA® TensorRT™ is an SDK for optimizing trained deep learning models to enable high-performance inference. TensorRT contains a deep learning inference optimizer for trained deep learning models, and a runtime for execution.

简单来说tensorrt是一个高性能推理部署框架，包括两个主要部分：优化器、执行器。

环境搭建

环境搭建主要包括三个部分：tensorrt， cuda， cudnn
建议使用docker，如百度paddle项目提供的gpu docker，里面已经安装好了cuda和cudnn，然后再下载一个 tensorrt即可。
docker下载：nvidia-docker pull registry.baidubce.com/paddlepaddle/paddle:2.2.2-gpu-cuda11.2-cudnn8
参考链接：百度paddlepaddle框架

tensorrt下载：
进入官网下载压缩包 官网下载链接

运行docker后，在自己的目录下下载Tensorrt安装包，解压，里面有头文件、库、可执行命令以及各种demo等。
cuda可以简单理解为一个工作台，一个编译器，一种语言。
cudnn是用cuda定义的专用的深度学习加速库。
如果把cuda类比成c++语言、 g++编译器， 那cudnn就是用c++写的一个动态库。
paddle docker环境下：cuda的已经安装好了，目录是：/usr/local/cuda/，里面有所需要的头文件和库。 cudnn动态库的目录是：/usr/lib/x86_64-linux-gnu/libcudnn.so

开始使用tensorrt

假设有一个很简单的模型，模型只有唯一的一层conv2d ：kernel-size=2*2;stride=1；假设输入是shape为[1,1,4,4]的全1数据，那么输出就是shape为[1,1,3,3]的tensor，且数值都为4。

完整的demo代码如下：

#include "NvInfer.h"
#include <iostream>
#include <cuda_runtime_api.h>
#include <vector>

class Logger : public nvinfer1::ILogger
{
public:
    void log(Severity severity, const char* msg) noexcept override
    {
        // suppress info-level messages
        if (severity <= Severity::kVERBOSE)
            std::cout << msg << std::endl;
    }
};

int main() {
  Logger logger;

  // Create an instance of the builder:
  nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);

  // Create a Network Definition
  uint32_t flag = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
  nvinfer1::INetworkDefinition* network = builder->createNetworkV2(flag);

  // Add the Input layer to the network
  auto input_data = network->addInput("input", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 4, 4});
  
  // Add the Convolution layer with hidden layer input nodes, strides and weights for filter and bias.
  std::vector<float>filter(2*2, 1.0);
  nvinfer1::Weights filter_w{nvinfer1::DataType::kFLOAT, filter.data(), 4};
  nvinfer1::Weights bias_w{nvinfer1::DataType::kFLOAT, nullptr, 0};
  auto conv2d = network->addConvolution(
                *input_data, 1, nvinfer1::DimsHW{2, 2}, filter_w, bias_w);
  conv2d->setStride(nvinfer1::DimsHW{1, 1});

  // Add a name for the output of the conv2d layer so that the tensor can be bound to a memory buffer at inference time:
  conv2d->getOutput(0)->setName("output");

  // Mark it as the output of the entire network:
  network->markOutput(*conv2d->getOutput(0));

  // Building an Engine(optimize the network)
  nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
  nvinfer1::IHostMemory*  serializedModel = builder->buildSerializedNetwork(*network, *config);
  
  nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
  nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(serializedModel->data(), serializedModel->size());

  // Prepare input_data
  int32_t inputIndex = engine->getBindingIndex("input");
  int32_t outputIndex = engine->getBindingIndex("output");
  std::vector<float> input(4*4, 1.0);
  std::vector<float> output(3*3);
  void *GPU_input_Buffer_ptr;  // a host ptr point to a GPU buffer
  void *GPU_output_Buffer_ptr;  // a host ptr point to a GPU buffer
  void* buffers[2];
  cudaMalloc(&GPU_input_Buffer_ptr, sizeof(float)*4*4); //malloc gpu buffer for input
  cudaMalloc(&GPU_output_Buffer_ptr, sizeof(float)*3*3); //malloc gpu buffer for output

  cudaMemcpy(GPU_input_Buffer_ptr, input.data(), input.size()*sizeof(float), cudaMemcpyHostToDevice); // copy input data from cpu to gpu


  buffers[inputIndex] = static_cast<void*>(GPU_input_Buffer_ptr);
  buffers[outputIndex] = static_cast<void*>(GPU_output_Buffer_ptr);

  // Performing Inference
  nvinfer1::IExecutionContext *context = engine->createExecutionContext();
  context->executeV2(buffers);

  // copy result data from gpu to cpu
  cudaMemcpy(output.data(), GPU_output_Buffer_ptr, output.size()*sizeof(float), cudaMemcpyDeviceToHost); 

  // display output
  std::cout << "output is : \n";
  for(auto i : output)
    std::cout << i << " ";
  std::cout << std::endl;
}

本demo涉及的api来自参考文档：
开发者文档第6.4.1节 和 第三节 The C++ API
代码很简单清晰，官方文档较完善，这里不再多述。
唯一一点是，代码中使用了cuda的两个api(位于头文件#include <cuda_runtime_api.h>)，cudaMemcpy：申请cpu内存，cudaMemcpy：数据在gpu和cpu之间拷贝。注意cudaMemcpy的参数，由于申请了一段gpu的内存之后，需要使用一个host的指针指向这个内存，所以这个api需要改动这个指针的指向，故传入这个指针的地址（二级指针）
cud api参考文档
demo中使用的其他api位于头文件NvInfer.h中，在下载的TensorRT-8.2.3.0安装包中。

编译

cuda所在目录：/usr/local/cuda/
cudnn库路径：/usr/lib/x86_64-linux-gnu/libcudnn.so
下载的TensorRT-8.2.3.0根目录： /TensorRT-8.2.3.0

编译指令如下：

g++ -std=c++11 /weishengying/main.cc  \
-I  /TensorRT-8.2.3.0/include/  \
-I  /usr/local/cuda/include  \
-L /usr/local/cuda/lib64 -lcudart  \
-L /usr/lib/x86_64-linux-gnu/ -lcudnn  \
-L /TensorRT-8.2.3.0/lib/ -lnvinfer

运行

./a.out
最终结果为：
output is : 
4 4 4 4 4 4 4 4 4 

其他参考

TensorRT github 仓库