zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

onnx_export.py

This commit is contained in:
Glenn Jocher 2020-06-01 14:05:01 -07:00
parent df7988d8d0
commit 7a00a5eff2
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
models/onnx_export.py
View File

@ -1,3 +1,7 @@
# Exports a pytorch *.pt model to *.onnx format. Example usage:
# $ export PYTHONPATH="$PWD"
# $ python models/onnx_export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
import argparse import argparse
import onnx import onnx
@ -6,7 +10,7 @@ from models.common import *
if __name__ == '__main__': if __name__ == '__main__':
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--weights', default='../weights/yolov5s.pt', help='model path RELATIVE to ./models/') parser.add_argument('--weights', default='./weights/yolov5s.pt', help='model path RELATIVE to ./models/')
parser.add_argument('--img-size', default=640, help='inference size (pixels)') parser.add_argument('--img-size', default=640, help='inference size (pixels)')
parser.add_argument('--batch-size', default=1, help='batch size') parser.add_argument('--batch-size', default=1, help='batch size')
opt = parser.parse_args() opt = parser.parse_args()