YOLO v5 pt模型转onnx和tensort

pt模型转onnx

• yolov5自带转onnx

python models/export.py --weights runs/exp12/weights/best.pt --img 640 --batch 1

• 安装onnx模型简化库onnx-simplifier

pip install onnx-simplifier

• 简化onnx模型

python -m onnxsim best.onnx best-sim.onnx

pt模型转tensorRT

• 转换 .wts文件

python gen_wts.py

将/tensorrtx/yolov5/gen_wts.py这个文件放到yolov5的根目录下。修改gen_wts.py文件中模型文件的目录为自己训练好的模型文件路径，会在模型文件根目录下生成yolov5s.wts文件。

• .wts文件转.engine文件（tensorrt模型）

static constexpr int CLASS_NUM = your_classes;

根据自己的类别修改yololayer.h第11行的CLASS_NUM参数

• 编译yolov5

cd yolov5
mkdir build
cd build
cmake …
make

• 转换

./yolov5 -s

将之前转好的wts模型复制到/tensorrtx/yolov5目录下，运行转换代码
注意：该步骤的转换最好在要部署机器上进行

• 验证C++接口

./yolov5 -d …/samples

进入build目录，运行转换代码；测试的图片放置在samples目录

参考

https://blog.csdn.net/qq_45057749/article/details/115016683