glenn-jocher commited on
Commit
264be1a
1 Parent(s): 2e53844

Rename `opset_version` to `opset` (#4135)

Browse files
Files changed (1) hide show
  1. export.py +5 -5
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, opset_version, train, dynamic, simplify):
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=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
- opset_version=12, # ONNX: opset version
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, opset_version, train, dynamic, simplify)
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-version', type=int, default=12, help='ONNX: opset version')
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