Yolov5——使用export.py把pt格式模型输出为onnx模型

前提依赖

需要Python>=3.8.0，并且含有PyTorch>=1.8，安装好yolov5所需要的所有依赖

git clone https://github.com/ultralytics/yolov5
# 用于克隆yolov5项目
cd yolov5
pip install -r requirements.txt  
# 安装所有依赖库

可转换的模型

Format	export.py --include	Model
PyTorch	-	yolov5s.pt
TorchScript	torchscript	yolov5s.torchscript
ONNX	onnx	yolov5s.onnx
OpenVINO	openvino	yolov5s_openvino_model/
TensorRT	engine	yolov5s.engine
CoreML	coreml	yolov5s.mlmodel
TensorFlow SavedModel	saved_model	yolov5s_saved_model/
TensorFlow GraphDef	pb	yolov5s.pb
TensorFlow Lite	tflite	yolov5s.tflite
TensorFlow Edge TPU	edgetpu	yolov5s_edgetpu.tflite
TensorFlow.js	tfjs	yolov5s_web_model/
PaddlePaddle	paddle	yolov5s_paddle_model/

export可选参数

可配置项	参数	解释
"--data"	type=str default=ROOT/"data/coco128.yaml" help="dataset.yaml path"	数据集配置目录，设置为训练模型时的数据集配置即可
"--weights"	nargs="+" type=str default=ROOT / "yolov5s.pt" help="model.pt path(s)"	权重文件目录
"--imgsz" "--img" "--img-size"	nargs="+" type=int default=[640, 640] help="image (h, w)"	输入模型的图片大小，默认为[640,640]
"--batch-size"	type=int default=1 help="batch size"	batch大小，每次处理的图片数，默认为1
"--device"	default="cpu" help="cuda device, i.e. 0 or 0,1,2,3 or cpu"	模型运行设备，默认为cpu
"--half"	action="store_true" help="FP16 half-precision export"	是否使用半精度FP16export转换，默认False
"--inplace"	action="store_true" help="set YOLOv5 Detect() inplace=True"	是否设置Detect()类的inplace属性为Treu，默认False
"--keras"	action="store_true" help="TF: use Keras"	是否在TensorFlow中使用Keras，默认False
"--optimize"	action="store_true" help="TorchScript: optimize for mobile"	是否对TorchScript进行移动端量化，默认False
"--int8"	action="store_true" help="CoreML/TF/OpenVINO INT8 quantization"	是否进行INT8量化，用于CoreML或TensorFlow，默认False
"--per-tensor"	action="store_true" help="TF per-tensor quantization"	是否对TensorFlow进行per-tensor量化，默认False
"--dynamic"	action="store_true" help="ONNX/TF/TensorRT: dynamic axes"	onnx转换参数，是否进行批处理变量，默认False
"--simplify"	action="store_true" help="ONNX: simplify model"	onnx模型转换参数,是否简化onnx模型，默认False
"--mlmodel"	action="store_true" help="CoreML: Export in *.mlmodel format"	是否将模型导出为CoreML格式，默认False
"--opset"	type=int default=17 help="ONNX: opset version"	onnx模型转换参数，设置onnx版本，默认17
"--verbose"	action="store_true" help="TensorRT: verbose log"	是否在 TensorRT 中提供详细日志，默认False
"--workspace"	type=int default=4 help="TensorRT: workspace size (GB)"	指定 TensorRT 的工作空间大小
"--nms"	action="store_true" help="TF: add NMS to model"	是否在 TensorFlow 中添加非极大值抑制（NMS），默认False
"--agnostic-nms"	action="store_true" help="TF: add agnostic NMS to model"	是否在 TensorFlow 中添加不可知的 NMS，默认False
"--topk-per-class"	type=int default=100 help="TF.js NMS: topk per class to keep"	在 TF.js 中，每个类别保留的检测框数量
"--topk-all"	type=int default=100 help="TF.js NMS: topk for all classes to keep"	在 TF.js 中，所有类别保留的检测框数量
"--iou-thres"	type=float default=0.45 help="TF.js NMS: IoU threshold"	在 TF.js 中，NMS 的 IoU 阈值，默认为0.45
"--conf-thres"	type=float default=0.25 help="TF.js NMS: confidence threshold"	在 TF.js 中，NMS 的置信度阈值，默认为0.25
"--include"	nargs="+" default=['torchscript', 'onnx'] help="torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle"	模型的输出格式，默认为[‘torchscript’,’onnx’]

nargs="+"

有这个参数表示用户需要至少提供一个输入文件，如果他们提供多个文件，这几个文件名会被放入一个列表

python script.py --input file1.txt file2.txt file3.txt
# args.input 将是一个列表，包含了所有传递给 --input 的值
# ['file1.txt', 'file2.txt', 'file3.txt']

action="store_true"

如果命令行中包含该选项，相应的变量将被设置为 True，否则为默认 False

如果修改默认为True（加入参数default=True）不管包不包含都会被设置为True

设置参数并运行

要把pt模型转换为onnx模型，需要指定pt模型位置，所使用的数据配置文件

# 在yolov5目录下,命令行运行
python export.py --weights yolov5s.pt --data data/coco.yaml --include onnx
# 一般指定pt模型位置，数据配置文件，要输出的模型格式这三个参数即可，其他参数可以自行增加

转换完成，可以使用detect.py进行执行检测，也可以使用val.py验证模型性能