Skip to content

Commit

Permalink
Scope onnx-simplifier requirements check (ultralytics#4730)
Browse files Browse the repository at this point in the history
* 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 <glenn.jocher@ultralytics.com>
  • Loading branch information
2 people authored and CesarBazanAV committed Sep 29, 2021
1 parent 665394f commit 5cfd471
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion export.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ def export_onnx(model, img, file, opset, train, dynamic, simplify):
# ONNX model export
prefix = colorstr('ONNX:')
try:
check_requirements(('onnx', 'onnx-simplifier'))
check_requirements(('onnx',))
import onnx

print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
Expand All @@ -66,6 +66,7 @@ def export_onnx(model, img, file, opset, train, dynamic, simplify):
# Simplify
if simplify:
try:
check_requirements(('onnx-simplifier',))
import onnxsim

print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
Expand Down

0 comments on commit 5cfd471

Please sign in to comment.