KaleiNeely glenn-jocher commited on
Commit
92e47b8
·
unverified ·
1 Parent(s): fbe67e4

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>

Files changed (2) hide show
  1. export.py +4 -5
  2. 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
- check_requirements(('onnx-simplifier',))
 
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 != 'cuda', '--optimize not compatible with cuda devices, i.e. use --device cpu'
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.3.6 # ONNX simplifier
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