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

结果

 

相关文章

     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')

 

笔者

笔者 - jxd