嵌入式模型部署学习笔记 ——在Jetson TX1上部署Yolov5模型

一、前言

前段时间在 Jetson TX1 上直接跑 YOLOv5 模型，由于 TX1 算力不够，实时检测画面出现卡帧现象，因为 YOLOv5 支持并且可以轻松地编译为 ONNX ，因此这里我们采用 TensorRT 对其进行推理加速处理。

二、移动端部署思路

YOLOv5 的部署，我们可以先在 x86 系统上将 Pytorch 的 .pt 文件转化成中间模型 .onnx 文件/ .wts 文件，然后在 TX1 上将中间模型 .onnx 文件/ .wts 文件转到 trt 的 engine ，最后通过 TensorRT 进行推理。

三、部署步骤

1. 克隆YOLOv5工程文件

git clone https://github.com/ultralytics/yolov5.git

2. pt 转 ONNX

安装 ONNX

在 YOLOv5 文件夹下运行

pip install -r requirements.txt onnx # install requirements.txt and ONNX

转换文件

YOLOv5 官方提供了将 .pt 模型移植到 TorchScript、ONNX、CoreML 格式的 .py 文件

我们只需要将 export.py 中加载 .pt 文件的路径修改为自己想要转换的模型路径即可，我这里使用 YOLOv5l 训练出来的 best.pt 文件进行操作

之后可以在 .pt 文件所在的文件夹中得到对应的 ONNX 文件

3. 移动端部署

接下来我们移步到 TX1 上进行进一步的部署

克隆工程

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

生成 .wts 文件

这一步可以在 x86 系统上操作
将 tensorrtx/yolov5/gen_wts.py文件 复制到 yolov5 文件夹中
在目录命令行下输入 .pt 文件路径即可生成对应的 .wts 文件

python3 gen_wts.py weights/best.pt

生成引擎文件

生成部署引擎，将 .wts 文件放到 tensorrtx/yolov5 文件夹中，可以在 yolov5.cpp 文件中修改 fp、int8、device、nms_thresh、conf_thresh、batch_size，这里我使用默认参数。

接下来我们在 yololayer.h 文件中修改模型类个数，我们还可以在 yololayer.h 文件中修改输入图片的尺寸，图片尺寸必须是32的倍数。缩小输入尺寸在一定程度上可以加快推理速度。

至此，我们已将 .wts 文件放到 tensorrtx/yolov5 文件夹中，并且在 yololayer.h 文件中修改了我们自己模型的类个数，接下来我们生成引擎文件。

在 tensorrtx/yolov5 文件夹中运行以下命令行

mkdir build
cd build
cmake ..
make -j6
sudo ./yolov5 -s ../best.wts best.engine l 
#sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw]         

到此，我们已经生成了 engine 部署引擎文件。

4. python 调用摄像头实时检测

在 TX1 上安装 pytorch 和 torchvision 比较麻烦，因此这里我们这里跳过了安装库的步骤直接对 yolov5_trt.py 文件进行修改。

注释掉 yolov5_trt.py 文件中的 import torch 和 import torchvision ，添加 import argparse

这一行直接改为 y = np.zeros_like(x)，创