zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

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

* 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>
This commit is contained in:
Zegorax 2021-09-09 16:49:10 +02:00 committed by GitHub
parent 1cad0ce2c7
commit 2d9411dbb8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
export.py
View File

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