結果
相關文章
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