tensorrt 推理yolov5-6.2模型

tensorrt 相关知识

tensorrt 是nvidia 公司推出的一款能够加速深度学习网络推理的SDK。

具体怎样加速的，什么类型的网络结构能够起到加速的作用等等不做赘述。

参考的版本来自https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5

实现步骤：

1.根据实际业务需求，利用pytorch yolov5-6.2 版本进行模型训练，网址链接为：https://github.com/ultralytics/yolov5得到.pt 文件。

2. .pt 文件转换成 .wts 文件。通过  运行 这个脚本gen_wts.py 。

3.编译运行yolov5 文件夹下的源代码，生成可执行文件。终端运行命令行demo 跑通：

sudo ./yolov5 -s yolov5s.wts yolov5s.engine s
sudo ./yolov5 -d yolov5s.engine ../images
 

关键代码解读，顺序是按照tensorrt 的执行流程展开：

1.cudaSetDevice(DEVICE);//设置GPU id ，决定将当前程序跑在那块GPU下。

2.nvinfer1::IHostMemory *modelStream{nullptr}; //IHostMemory 是经tensorrt 优化后的序列化模型被保存的对象。

3.nvinfer1::IBuilder *builder = nvinfer1::createInferBuilder(gLogger);//想要经过tensorrt 序列化模型，首先需要创建一个builder 实例。

4.nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();//建一个IBuilderConfig对象来告诉TensorRT该如何对模型进行优化。这个接口定义了很多属性，其中最重要的一个属性是工作空间的最大容量。在网络层实现过程中通常会需要一些临时的工作空间，这个属性会限制最大能申请的工作空间的容量，如果容量不够的话会导致该网络层不能成功实现而导致错误。另外，还可以通过这个对象设置模型的数据精度。

5.ICudaEngine *engine = nullptr; //实例化一个构建本地保存序列化后的模型引擎对象。

6.INetworkDefinition* network = builder->createNetworkV2(0U); //模型序列化后要构建一个模型网络结构。

7.DataType::kFLOAT //决定模型序列化后的精度。

kFLOAT //!< FP32 format.
kHALF  //!< FP16 format.
kINT8  //!< INT8 format. //INT8模式只支持计算能力为6.1以上的GPU
kINT32 //!< INT32 format.
kTF32  //!< TF32 format.

8.engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name);

//通过这个函数创建engine 文件。

9.(*modelStream) = engine->serialize(); engine 文件进行模型序列化。

以上完成tensorrt 序列化模型文件的过程。

实际使用过程，本地直接会生成tensorrt序列化后的模型文件。

1.IRuntime* runtime = createInferRuntime(gLogger); //创建一个IRuntime对象。

2.ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);//反序列化engine 文件。

以上不论是通过构建模型的方式，还是直接加载本地文件进行反序列化，最终完成engine 的创建。

模型推理流程：

IExecutionContext* context = engine->createExecutionContext();//ICudaEngine对象中存放着经过TensorRT优化后的模型，不过如果要用模型进行推理则还需要通过createExecutionContext()函数去创建一个IExecutionContext对象来管理推理的过程：

cudaStream_t stream;
cudaStreamCreate(&stream); //使用下面的方法可以声明和创建一个显式流： 目的就是提什程序的执行效率。

CUDA_CHECK(cudaMallocHost((void**)&img_host, MAX_IMAGE_INPUT_SIZE_THRESH * 3));

// prepare input data cache in pinned memory

CUDA_CHECK(cudaMalloc((void**)&img_device, MAX_IMAGE_INPUT_SIZE_THRESH * 3));

// prepare input data cache in device memory

memcpy(img_host, in.data, size_image);

//copy data to pinned memory

//copy data to device memory

CUDA_CHECK(cudaMemcpyAsync(img_device, img_host, size_image, cudaMemcpyHostToDevice, stream)); // 将数据从本地拷贝到设备上。

context.enqueue(batchSize, buffers, stream, nullptr);//模型进行推理。

CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); //将推理后的结果从设备拷贝到本地。

// 清理CUDA流

cudaStreamSynchronize(stream);

cudaStreamDestroy(stream);

模型推理结束

之后就是后处理，不做详解。