Zegorax glenn-jocher commited on
Commit
2d9411d
·
unverified ·
1 Parent(s): 1cad0ce

Scope `onnx-simplifier` requirements check (#4730)

Browse files

* Changed onnx-simplifier check behavior

Export.py has been updated to check for onnx-simplifier requirement only when the --simplify argument is added.
Allows for better flexibility and one less requirement if simplify is not needed.

* Fix single-element tuples

Co-authored-by: Glenn Jocher <[email protected]>

Files changed (1) hide show
  1. export.py +2 -1
export.py CHANGED
@@ -44,7 +44,7 @@ def export_onnx(model, img, file, opset, train, dynamic, simplify):
44
  # ONNX model export
45
  prefix = colorstr('ONNX:')
46
  try:
47
- check_requirements(('onnx', 'onnx-simplifier'))
48
  import onnx
49
 
50
  print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
@@ -66,6 +66,7 @@ def export_onnx(model, img, file, opset, train, dynamic, simplify):
66
  # Simplify
67
  if simplify:
68
  try:
 
69
  import onnxsim
70
 
71
  print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
 
44
  # ONNX model export
45
  prefix = colorstr('ONNX:')
46
  try:
47
+ check_requirements(('onnx',))
48
  import onnx
49
 
50
  print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
 
66
  # Simplify
67
  if simplify:
68
  try:
69
+ check_requirements(('onnx-simplifier',))
70
  import onnxsim
71
 
72
  print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')