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的方案
先转换为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<std::string, Weights> 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的编写的大佬随意哈)