使用TensorRT加速WongKinYiu/Pytorch_YOLOv4-tiny

HFUT_Xin

已于 2022-12-03 11:08:45 修改

阅读量577

点赞数

文章标签： pytorch python 目标检测

于 2022-11-20 22:16:02 首次发布

本文链接：https://blog.csdn.net/DGHAHSFIA/article/details/127954770

版权

在此对基于Tensorrtx 实现 darknet Yolov4-tiny 模型、tjuskyzhang/Scaled-YOLOv4-TensorRT的作者表示感谢，本文建立在上述文章之上。

文章目录

1、环境和使用的YOLOv4-tiny源码

环境：Jetson TX2 , Tensorrt 7.1
YOLOv4-tiny源码： https://github.com/WongKinYiu/PyTorch_YOLOv4/tree/master，branch master

2、使用 cfg 和 pt 生成 wts

将 gen_wts.py 复制到 ScaledYOLOv4/ 下
	git clone https://github.com/wang-xinyu/tensorrtx.git 	
	git clone -b yolov4-csp https://github.com/WongKinYiu/ScaledYOLOv4.git
	cp  {tensorrtx}/scaled-yolov4/gen_wts.py  {ScaledYOLOv4/} 	
	cd  {ScaledYOLOv4/} 
生成 wts 	
	//将要转换的 cfg、pt 文件复制到 ScaledYOLOv4/ 路径下
	python  gen_wts.py  best_overall.pt  	
	// 执行前先修改一下 gen_wts.py 文件内容，执行后 a file 'yolov4-tiny.wts' will be generated.

注意核对一下 cfg 文件路径和图片输入尺寸
gen_wts.py

3、使用 wts 生成 engine，然后推理加速

git clone https://github.com/tjuskyzhang/Scaled-YOLOv4-TensorRT.git 
cp  {ScaledYOLOv4}/yolov4-tiny.wts  {Scaled-YOLOv4-TensorRT}/yolov4-tiny-tensorrt 
cd  Scaled-YOLOv4-TensorRT/yolov4-tiny-tensorrt

修改 yololayer.h：
打开 Scaled-YOLOv4-TensorRT/yolov4-tiny-tensorrt 路径下的 yololayer.h ，修改 yololayer.h ，需要修改的地方已在下面图片标注。注意：yolo2 下面使用的 anchors 是 anchors 1 2 3，不是 anchors 0 1 2。具体原因见解释说明。

解释说明

来源：https://github.com/hunglc007/tensorflow-yolov4-tflite/issues/111#issuecomment-650205659

修改 yololayer.cu：
将 yololayer.cu 中188～191改成：

det->bbox[0] = (col + Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid]) * 2 -0.5) * INPUT_W / yoloWidth;
det->bbox[1] = (row + Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid]) * 2 -0.5) * INPUT_H / yoloHeight;
det->bbox[2] = (2 * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid])) * (2 * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid])) * anchors[2*k];
det->bbox[3] = (2 * Logist(curInput[idx + k * info_len_i * total_grid + 3 * total_grid])) * (2 * Logist(curInput[idx + k * info_len_i * total_grid + 3 * total_grid])) * anchors[2*k + 1];

来源：issues/29#

然后继续执行命令

mkdir build 
cd  build 
cmake  .. 
make  

// Serialize the model and generate yolov4-tiny.engine 
./yolov4-tiny  -s      // 执行窗口出现 set name out 时表示正在转换，稍等一会 
// Deserialize and generate the detection results _dog.jpg and so on. 
./yolov4-tiny  -d   the/path/of/your/images