Upgrade onnxsim to v0.4.1 (#8632)
Browse files* upgrade onnxsim to v0.4.1
Signed-off-by: daquexian <daquexian566@gmail.com>
* Update export.py
* Update export.py
* Update export.py
* Update export.py
* Update export.py
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
- export.py +4 -5
- requirements.txt +1 -1
export.py
CHANGED
@@ -152,13 +152,12 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
|
|
152 |
# Simplify
|
153 |
if simplify:
|
154 |
try:
|
155 |
-
|
|
|
156 |
import onnxsim
|
157 |
|
158 |
LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
|
159 |
-
model_onnx, check = onnxsim.simplify(model_onnx
|
160 |
-
dynamic_input_shape=dynamic,
|
161 |
-
input_shapes={'images': list(im.shape)} if dynamic else None)
|
162 |
assert check, 'assert check failed'
|
163 |
onnx.save(model_onnx, f)
|
164 |
except Exception as e:
|
@@ -493,7 +492,7 @@ def run(
|
|
493 |
imgsz *= 2 if len(imgsz) == 1 else 1 # expand
|
494 |
assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
|
495 |
if optimize:
|
496 |
-
assert device.type
|
497 |
|
498 |
# Input
|
499 |
gs = int(max(model.stride)) # grid size (max stride)
|
|
|
152 |
# Simplify
|
153 |
if simplify:
|
154 |
try:
|
155 |
+
cuda = torch.cuda.is_available()
|
156 |
+
check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
|
157 |
import onnxsim
|
158 |
|
159 |
LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
|
160 |
+
model_onnx, check = onnxsim.simplify(model_onnx)
|
|
|
|
|
161 |
assert check, 'assert check failed'
|
162 |
onnx.save(model_onnx, f)
|
163 |
except Exception as e:
|
|
|
492 |
imgsz *= 2 if len(imgsz) == 1 else 1 # expand
|
493 |
assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
|
494 |
if optimize:
|
495 |
+
assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'
|
496 |
|
497 |
# Input
|
498 |
gs = int(max(model.stride)) # grid size (max stride)
|
requirements.txt
CHANGED
@@ -25,7 +25,7 @@ seaborn>=0.11.0
|
|
25 |
# Export --------------------------------------
|
26 |
# coremltools>=4.1 # CoreML export
|
27 |
# onnx>=1.9.0 # ONNX export
|
28 |
-
# onnx-simplifier>=0.
|
29 |
# nvidia-pyindex # TensorRT export
|
30 |
# nvidia-tensorrt # TensorRT export
|
31 |
# scikit-learn==0.19.2 # CoreML quantization
|
|
|
25 |
# Export --------------------------------------
|
26 |
# coremltools>=4.1 # CoreML export
|
27 |
# onnx>=1.9.0 # ONNX export
|
28 |
+
# onnx-simplifier>=0.4.1 # ONNX simplifier
|
29 |
# nvidia-pyindex # TensorRT export
|
30 |
# nvidia-tensorrt # TensorRT export
|
31 |
# scikit-learn==0.19.2 # CoreML quantization
|