Jetson nano上使用trt部署自己训练的YOLOv5模型（亲测可用！！）

yolov5的github链接:GitHub - ultralytics/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLite

TensorRT下载链接:tensorrtx/yolov5 at master · wang-xinyu/tensorrtx · GitHub

---

查看自己使用的yolov5版本

以yolov5s为例，打开文件yolo-master/models/yolov5s.yaml查看yolov5版本

---

下载trt并导出wts格式模型文件

链接tensorrtx/yolov5 at master · wang-xinyu/tensorrtx · GitHub可查看yolov5对应版本的trt并下载

下载trt得到tensorrtx文件夹，将tensorrtx/yolov5/gen_wts.py拷贝至yolov5项目根目录下

执行命令生成wts文件

python gen_wts.py yolov5/runs/train/exp4/weights/best.pt

运行完毕之后会生成一个.wts文件放到tensorrtx/yolov5/目录下

---

Ubantu部署

在ubantu上也克隆一个tensorrtx项目

将生成的.wts拷贝到tensorrtx/yolov5/下

修改tensorrtx/yolov5/yololayer.h中

static constexpr int CLASS_NUM = 4;  （我自己的类别为4）

在tensorrtx/yolov5/目录下打开终端，编译代码

mkdir build
cd build
cmake ..
make

将.wts文件转为.engine文件

./yolov5 -s ../best.wts ../best.engine s    (n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw)

执行推理任务

./yolov5 -d ../best.engine ../images

执行完在build目录生成检测完成的图片

使用python推理，编辑tensorrtx\yolov5中的yolov5_trt.py或者yolov5_trt_cuda_python.py使用python推理。

若能帮到您的话，点赞评论吧！！