glenn-jocher
commited on
Commit
•
7a00a5e
1
Parent(s):
df7988d
onnx_export.py
Browse files- models/onnx_export.py +5 -1
models/onnx_export.py
CHANGED
@@ -1,3 +1,7 @@
|
|
|
|
|
|
|
|
|
|
1 |
import argparse
|
2 |
|
3 |
import onnx
|
@@ -6,7 +10,7 @@ from models.common import *
|
|
6 |
|
7 |
if __name__ == '__main__':
|
8 |
parser = argparse.ArgumentParser()
|
9 |
-
parser.add_argument('--weights', default='
|
10 |
parser.add_argument('--img-size', default=640, help='inference size (pixels)')
|
11 |
parser.add_argument('--batch-size', default=1, help='batch size')
|
12 |
opt = parser.parse_args()
|
|
|
1 |
+
# Exports a pytorch *.pt model to *.onnx format. Example usage:
|
2 |
+
# $ export PYTHONPATH="$PWD"
|
3 |
+
# $ python models/onnx_export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
|
4 |
+
|
5 |
import argparse
|
6 |
|
7 |
import onnx
|
|
|
10 |
|
11 |
if __name__ == '__main__':
|
12 |
parser = argparse.ArgumentParser()
|
13 |
+
parser.add_argument('--weights', default='./weights/yolov5s.pt', help='model path RELATIVE to ./models/')
|
14 |
parser.add_argument('--img-size', default=640, help='inference size (pixels)')
|
15 |
parser.add_argument('--batch-size', default=1, help='batch size')
|
16 |
opt = parser.parse_args()
|