使用tensorRT部署yolov5 v5.0—windows

目录

• 前言
• step 1. 下载yolov5源码以及tensorrtx源码，并将yolov5s.pt转为.wts模型
• step 2. 安装cuDNN、TensorRT以及验证TensorRT是否安装成功
• step 3. https://github.com/wang-xinyu 使用vs跑通github中yolov5 5.0 的C++代码
• step 4. 跑自己训练的模型
• 总结

前言

    前段时间研究了Pytorch的环境配置，之后便从github上下载了yolov5的源码，并在自己的电脑端配置好对应的环境并运行，最后发现生成的权重文件yolov5s.pt不仅可以通过量化压缩成onxx模型，而且还可以使用TensorRT推理加速生成engine模型，这对使得模型部署在移动端具有很大的优势，于是便尝试着在自己的电脑上通过TensorRT部署yolov5模型。
     现在网上有很多可以参考的博客，但大多数都是针对某一个环节进行了仔细的解释说明，这在前期的学习中不免会让人产生云里雾里的感觉，难以从一个全局的角度去看待这个问题，换句话说就是很少有把整个流程先总结下来，先让我们知道需要那些模块，该准备些什么模块，以及这些模块之间又有什么样的联系，然后再细分到各个小模块去说明解释。所以今天就从这个角度去发出，总结一下最近学习的一些内容。
在此之前假设你已经掌握了Pytorch、CUDA、cuDNN的基础知识以及配置好了yolov5的环境并调试运行过源码

注 :如果还没有掌握上述基础知识，可以参考另外两篇博文
1.深度学习之Pytorch环境搭建
2.yolov5部署之环境配置及源码测试

首先整个过程分为以下7个步骤（模块）：

step 1. 下载yolov5源码以及tensorrtx源码，并将yolov5s.pt转为.wts模型。
step 2. 安装cuDNN、TensorRT以及验证TensorRT是否安装成功
step 3. https://github.com/wang-xinyu 使用vs跑通github中yolov5 5.0 的C++代码
step 4. 跑自己训练的模型

下面将从这个7个步骤，逐步进行解释，首先贴出我的环境：

系统 : WIN 10
cuda_10.1
cudnn-10.1-windows10-x64-v7.6.5.32
TensorRT-6.0.1.5.Windows10.x86_64.cuda-10.1.cudnn7.6
OpenCV 4.4.0
vs2015

• https://developer.nvidia.com/rdp/cudnn-download
  https://developer.nvidia.com/tensorrt

  代码版本：

  yolov5 5.0
  tensorrtx 5.0

git clone -b v5.0 https://github.com/ultralytics/yolov5.git
git clone -b yolov5-v5.0 https://github.com/wang-xinyu/tensorrtx.git

这里不得不说一下wang-xinyu的工作实在是太赞了！ 源码地址点这里

step 1. 下载yolov5源码以及tensorrtx源码，并将yolov5s.pt转为.wts模型

• 将tensorrtx源码中的gen_wts.py复制到yolov5源码中并运行，生成.wts模型。

step 2. 安装cuDNN、TensorRT以及验证TensorRT是否安装成功

1. 安装cuDNN

1 将cuDNN压缩包解压
2 将cuda\bin中的文件复制到 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\bin
3 将cuda\include中的文件复制到 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\include
4 将cuda\lib中的文件复制到 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\lib

4步完成cuDNN安装

2. 安装TensorRT

1 将TensorRT压缩包解压
2 将 TensorRT-6.0\include中头文件复制到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\include
3 将TensorRT-6.0\lib中所有lib文件复制到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\lib\x64
4 将TensorRT-6.0\lib中所有dll文件复制到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\bin

4步完成TensorRT安装

3. 验证TensorRT是否安装成功

1-用VS2015打开 TensorRT-6.0\samples\sampleMNIST\sample_mnist.sln
2- 在VS2015中，右键工程，选择属性->配置属性->常规->目标平台版本->8.1
3- 在VS2015中，右键工程，选择属性->配置属性->常规->平台工具集->vs2015(v140)
4- 右键工程->重新生成
5- 用anaconda 进入TensorRT-6.0\data\mnist目录，运行python download_pgms.py
6- 进入TensorRT-6.0\bin，双击sample_mnist.exe，如果没有报错则说配置成功。

 

5. 打开yolov5的VS工程，生成engine模型

1- 在刚刚设置的build the binaries路径下，打开yolov5的工程
2- 编译生成
3- 可以看到在\tensorrtx\yolov5\build\Debug文件下生成了一个yolov5.exe文件
4- cmd进入到\tensorrtx\yolov5\build\Debug目录下，然后执行yolov5.exe -s命令, 可以看到在当前目录下生成了一个yolov5.engine文件，说明转换成功。
5- 在\tensorrtx\yolov5\build\Debug目录下新建一个samples文件夹然后放入测试图片，最好通过cmd执行yolov5.exe -d ../samples，调用yolov5.engine进行测试

step 3. https://github.com/wang-xinyu 使用vs跑通github中yolov5 5.0 的C++代码

1. 将yolov5 v5.0的.h/.cpp/.cu代码，拷贝到已配好opencv环境的vs工程

2.配置环境变量

包含目录
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\include;$(SolutionDir)\opencv4.4\include;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\bin;$(IncludePath)

库目录
$(SolutionDir)\opencv4.4\lib;$(SolutionDir)\x64\Release;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.1\lib\x64;$(LibraryPath)

附加依赖项
cublas.lib
cudnn.lib
nvinfer.lib
nvinfer_plugin.lib
nvonnxparser.lib
nvparsers.lib

3、更改代码

将utils.h中的#include 注释掉，读图用glob读

	std::vector file_names;
	cv::glob("C:\\Users\\Ring\\Desktop\\bushu 3.0\\images", file_names);
	if (file_names.size() < 0) {
		std::cout << "read_files_in_dir failed." << std::endl;
		return -1;
	}

4，将main中的代码改一下，改到程序执行

cudaSetDevice(DEVICE);

    std::string wts_name = "../best.wts";
    //std::string engine_name = "";
    bool is_p6 = false;
    float gd = 0.33f, gw = 0.5f;
    std::string img_dir;
    /*if (!parse_args(argc, argv, wts_name, engine_name, is_p6, gd, gw, img_dir)) {
        std::cerr << "arguments not right!" << std::endl;
        std::cerr << "./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw]  // serialize model to plan file" << std::endl;
        std::cerr << "./yolov5 -d [.engine] ../samples  // deserialize plan file and run inference" << std::endl;
        return -1;
    }*/
	char *trtModelStream{ nullptr };
	size_t size{ 0 };
	std::string engine_name = STR2(NET);
	engine_name = "best" + engine_name + ".engine";
    // create a model using the API directly and serialize it to a stream
    if (false) {
        IHostMemory* modelStream{ nullptr };
        APIToModel(BATCH_SIZE, &modelStream, is_p6, gd, gw, wts_name);
        assert(modelStream != nullptr);
        std::ofstream p(engine_name, std::ios::binary);
        if (!p) {
            std::cerr << "could not open plan output file" << std::endl;
            return -1;
        }
        p.write(reinterpret_cast(modelStream->data()), modelStream->size());
        modelStream->destroy();
        return 0;
    }
	else {
		std::ifstream file(engine_name, std::ios::binary);
		if (file.good()) {
			file.seekg(0, file.end);
			size = file.tellg();
			file.seekg(0, file.beg);
			trtModelStream = new char[size];
			assert(trtModelStream);
			file.read(trtModelStream, size);
			file.close();
		}
	}

    // deserialize the .engine and run inference
    std::ifstream file(engine_name, std::ios::binary);
    if (!file.good()) {
        std::cerr << "read " << engine_name << " error!" << std::endl;
        return -1;
    }

    //char *trtModelStream = nullptr;
    //size_t size = 0;
    file.seekg(0, file.end);
    size = file.tellg();
    file.seekg(0, file.beg);
    trtModelStream = new char[size];
    assert(trtModelStream);
    file.read(trtModelStream, size);
    file.close();

step 3. 跑自己的模型

1. 用yolov5训练出自己的模型，生成.pt

2.将tensorrtx源码中的gen_wts.py复制到yolov5源码中并运行，把pt生成.wts模型。

3、将.wts放在该工程外面，在C++中将if(true)，生成.engine文件。

4，将if(true)改为if(flase)，使用.engine文件，开始跑吧！