glenn-jocher
commited on
Rename `opset_version` to `opset` (#4135)
Browse files
export.py
CHANGED
@@ -38,7 +38,7 @@ def export_torchscript(model, img, file, optimize):
|
|
38 |
print(f'{prefix} export failure: {e}')
|
39 |
|
40 |
|
41 |
-
def export_onnx(model, img, file,
|
42 |
# ONNX model export
|
43 |
prefix = colorstr('ONNX:')
|
44 |
try:
|
@@ -47,7 +47,7 @@ def export_onnx(model, img, file, opset_version, train, dynamic, simplify):
|
|
47 |
|
48 |
print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
|
49 |
f = file.with_suffix('.onnx')
|
50 |
-
torch.onnx.export(model, img, f, verbose=False, opset_version=
|
51 |
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
|
52 |
do_constant_folding=not train,
|
53 |
input_names=['images'],
|
@@ -108,7 +108,7 @@ def run(weights='./yolov5s.pt', # weights path
|
|
108 |
optimize=False, # TorchScript: optimize for mobile
|
109 |
dynamic=False, # ONNX: dynamic axes
|
110 |
simplify=False, # ONNX: simplify model
|
111 |
-
|
112 |
):
|
113 |
t = time.time()
|
114 |
include = [x.lower() for x in include]
|
@@ -149,7 +149,7 @@ def run(weights='./yolov5s.pt', # weights path
|
|
149 |
if 'torchscript' in include:
|
150 |
export_torchscript(model, img, file, optimize)
|
151 |
if 'onnx' in include:
|
152 |
-
export_onnx(model, img, file,
|
153 |
if 'coreml' in include:
|
154 |
export_coreml(model, img, file)
|
155 |
|
@@ -170,7 +170,7 @@ def parse_opt():
|
|
170 |
parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
|
171 |
parser.add_argument('--dynamic', action='store_true', help='ONNX: dynamic axes')
|
172 |
parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
|
173 |
-
parser.add_argument('--opset
|
174 |
opt = parser.parse_args()
|
175 |
return opt
|
176 |
|
|
|
38 |
print(f'{prefix} export failure: {e}')
|
39 |
|
40 |
|
41 |
+
def export_onnx(model, img, file, opset, train, dynamic, simplify):
|
42 |
# ONNX model export
|
43 |
prefix = colorstr('ONNX:')
|
44 |
try:
|
|
|
47 |
|
48 |
print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
|
49 |
f = file.with_suffix('.onnx')
|
50 |
+
torch.onnx.export(model, img, f, verbose=False, opset_version=opset,
|
51 |
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
|
52 |
do_constant_folding=not train,
|
53 |
input_names=['images'],
|
|
|
108 |
optimize=False, # TorchScript: optimize for mobile
|
109 |
dynamic=False, # ONNX: dynamic axes
|
110 |
simplify=False, # ONNX: simplify model
|
111 |
+
opset=12, # ONNX: opset version
|
112 |
):
|
113 |
t = time.time()
|
114 |
include = [x.lower() for x in include]
|
|
|
149 |
if 'torchscript' in include:
|
150 |
export_torchscript(model, img, file, optimize)
|
151 |
if 'onnx' in include:
|
152 |
+
export_onnx(model, img, file, opset, train, dynamic, simplify)
|
153 |
if 'coreml' in include:
|
154 |
export_coreml(model, img, file)
|
155 |
|
|
|
170 |
parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
|
171 |
parser.add_argument('--dynamic', action='store_true', help='ONNX: dynamic axes')
|
172 |
parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
|
173 |
+
parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
|
174 |
opt = parser.parse_args()
|
175 |
return opt
|
176 |
|