如何将Yolov4自定义训练权重转换为ONNX格式以用于TensorRT推理？

一、备齐必要文件和环境 #

• 得有训练好的best.weights、对应自定义网络的配置文件（比如yolov4-custom.cfg），还有写好类别名称的obj.names文件
• 装好依赖：Python 3.x，以及opencv-python、onnx这些库；要是用Darknet原生转换，还得编译Darknet的Python绑定

二、将Darknet权重转换为ONNX文件 #

这里提供两种实用方法：

方法1：用Darknet原生Python工具转换 #

• 先克隆Darknet仓库，按照仓库说明编译好Python绑定
• 写个简单的转换脚本，示例代码如下：

import darknet

# 替换成你的实际文件路径
cfg_path = "./yolov4-custom.cfg"
weights_path = "./best.weights"
names_path = "./obj.names"
output_onnx = "./yolov4-custom.onnx"

# 加载Darknet网络和元数据
net = darknet.load_net_custom(cfg_path.encode("ascii"), weights_path.encode("ascii"), 0, 1)
meta = darknet.load_meta(names_path.encode("ascii"))

# 导出ONNX文件
darknet.save_onnx(net, meta, output_onnx.encode("ascii"))

• 运行脚本，就能得到目标ONNX文件

方法2：用第三方工具（如ultralytics YOLO）转换 #

• 先安装ultralytics库：pip install ultralytics
• 直接用命令行转换，注意自定义Yolov4的配置文件参数要和ultralytics兼容：

yolo export model=yolov4-custom.cfg weights=best.weights format=onnx

三、验证ONNX文件有效性 #

• 用ONNX官方工具检查文件合法性：

onnxchecker yolov4-custom.onnx

• 也可以用netron可视化模型结构，确认没问题：

pip install netron
netron yolov4-custom.onnx

四、转换为TensorRT推理引擎 #

• 用TensorRT自带的trtexec工具生成推理引擎：

trtexec --onnx=yolov4-custom.onnx --saveEngine=yolov4-custom.engine --explicitBatch

• 如果需要FP16精度加速，添加--fp16参数：

trtexec --onnx=yolov4-custom.onnx --saveEngine=yolov4-custom-fp16.engine --explicitBatch --fp16

• 之后就可以用生成的.engine文件进行TensorRT推理了

内容的提问来源于stack exchange，提问作者Sidharth S.K