Tensorrt 7.2.3 转换yolov5 v3.0 pytorch1.8.1 onnx1.6.0

Tensorrt 转换yolov5

安装环境：

pytorch 1.8.1

torchvision 0.9.1

onnx 1.6.0

Tensorrt 7.2.3

yolov5 v3.0 l

权重为自己训练的6类

这里用了两种方案：

1.采用 wang xinyu的方案

github地址

先转换为wts，再进行编译，进行检测

具体流程：

git clone https://github.com/wang-xinyu/tensorrtx.git

###将 yolov5/gen_wts.py 复制到自己的yolov5文件夹下，修改gen_wts.py的路径等 注意device需设置为cpu，改成GPU可能有变量设备不一致的问题。
cd ···/yolov5-master
python gen_wts.py ###生成wts文件
cp best.wts ···/tensorrtx/yolov5/


cd ···/tensorrtx/yolov5/

下一步：
修改tensorrtx/yolov5/CMakeLists.txt

• 修改cuda库的路径

  cuda

	#cuda
    include_directories(/usr/local/cuda-10.2/targets/x86_64-linux/include)
    link_directories(/usr/local/cuda-10.2/targets/x86_64-linux/lib)

• 修改tensorrt的路径

  tensorrt

	#tensorrt
    include_directories(···/tensorrt/TensorRT-7.2.3.4/include/)
    link_directories(···/tensorrt/TensorRT-7.2.3.4/targets/x86_64-linux-gnu/lib/)

下一步：
需要对yololayer.h 和yolov5.cpp进行修改

• yololayer.h
  对第19行等进行修改

   static constexpr int CLASS_NUM = 6;   //根据权重类别修改
   static constexpr int INPUT_H = 608;   //根据权重修改图像resize尺寸
   static constexpr int INPUT_W = 608;

• 和yolov5.cpp
  对第7行等进行修改

   #define USE_FP16  // comment out this if want to use FP32
   #define DEVICE 1  // GPU id
   #define NMS_THRESH 0.4
   #define CONF_THRESH 0.1
   #define BATCH_SIZE 1

   #define NET l  // s m l x 这里我用 l 训练的

这里我用的是l,所以还需要对ICudaEngine* createEngine_l 进行修改，

std::map weightMap = loadWeights("../best.wts");//第215行改为
if (argc == 2 && std::string(argv[1]) == "-l")//第421行改为
std::cerr << "./yolov5 -l  // serialize model to plan file" << std::endl;//第446行改为

接下来就可以编译了

mkdir build && cd build
cmake -D CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda-10.2 .. 
make  //cmake失败了，make就没必要了，所以上一步必须搞定呀
./yolov5 -l   //会生成一个engine,基本成功了

接下来，两种方法推理

python yolov5_trt.py
or
./yolov5 -d  ···/images_folders

2.采用 pt—onnx—trt方案

(参考同事的代码，不在这里展示)
因为我用的是v3.0的export.py进行转换onnx，但是一直报错，然后看是旧版的yolov5的export有问题，所以我用v5.0的yolov5对权重进行onnx转换。

python export.py //生成best.onnx

根据需求修改gen_calibration_new.py，vehicle_infer.py

python gen_calibration_new.py //生成best.trt
python vehicle_infer.py //进行推理

3.两个方案的pros and cons

• 效果对比

  ** 第一种
  

  **第二种
  

• 性能对比

  该结果存疑，这里测试的有些问题，python调用转换成的engine或者trt时，都会突然出现显存占用变大，推理速度变慢的问题，emmm，有大佬能够指正一下为啥不胜感激~

  yolov5s模型	batch_size	显存	精度	推理速度
  pytorch	1	1025MB	fp32	52~57ms （python）
  tensorrt 1	1	897MB	fp32	39~45ms（可执行文件）
  tensorrt 2	1	1109MB	fp32	74ms(python)

• 难度

  感觉可控性perfer c++版本的，除了改c代码的难度，其他的c++代码的透明度更高一点，c++的可执行文件效率也会更高一点。
  第二种方式转换的时候，由于onnx的黑盒特性，转换为的模型在netron查看时，会出现tensor维度变化的问题，在后面推理时，感觉第一种的推理脚本自由度更高，可以自由修改（懂cuda的编写的大佬随意哈）