Python —— Win10将Yolov5的.pt模型导出为.onnx模型

原创已于 2023-02-27 18:29:56 修改 · 2.2k 阅读

48 ·

CC 4.0 BY-SA版权

automoblie0

文章标签：

#python #YOLO #onnx #pt

于 2023-02-24 15:31:51 首次发布

Python 专栏收录该内容

15 篇文章

订阅专栏

在Windows10环境下，使用conda激活Pytorch1.8.1环境，然后在Yolov5源代码目录下，修改export.py文件以适应自己的训练数据和权重文件。通过降低ONNX的opset版本解决导出失败的问题，并指定导出格式为ONNX。如果需要，可以修改设备参数以使用GPU导出模型。

部署运行你感兴趣的模型镜像

结果

在这里插入图片描述

Python —— Pytorch训练集指令、测试数据指令

Python —— Windows10下配置Pytorch环境、进行训练模型并测试（完整流程，附有视频）

正文

参考：如何在windows系统下将yolov5的pt模型导出为onnx模型

# 进入conda,进入自己的pytorch环境
conda activate pytorch1.8.1

# 进入yolov5源码目录
cd D:\Software\Anaconda\newCODE\yolov5

# 修改export.py
	将:
		parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
    改为自己训练时用的yaml:
    	parser.add_argument('--data', type=str, default=ROOT / 'data/voc-mask.yaml', help='dataset.yaml path')
    
    将:
    	parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
    改为自己训练产生的pt模型:
    	parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'runs/train/base/weights/best.pt', help='model.pt path(s)')

# 生成.onnx模型。导出的.onnx会创建在.pt的同级目录下
python export.py

# 附错误问题：
	1、执行生成onnx模型指令出现错误:
		ONNX: export failure 0.0s: Unsupported ONNX opset version: 17
		解决:
			修改export.py
				将:
					parser.add_argument('--opset', type=int, default=17, help='ONNX: opset version')
    			改为:
    				parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')

	2、如果导出没有onnx只有torchscript,则:
		将:
			parser.add_argument('--include',nargs='+',default=['torchscript'],help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle')
		改为:
			parser.add_argument('--include',nargs='+',default=['onnx'],help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle')


#附加知识:
	1、默认使用cpu导出.onnx模型,如需要用gpu导出:
		将:
			parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
		改为:
			parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')