天天看點

Python —— Win10将Yolov5的.pt模型導出為.onnx模型

結果
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