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

<