mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Upgrade onnxsim to v0.4.1 (#8632) 

* 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>
pull/8664/head
daquexian 2022-07-20 01:01:24 +08:00 committed by GitHub
parent fbe67e4653
commit 92e47b85d9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 5 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
export.py
View File

@ -152,13 +152,12 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
# Simplify # Simplify
if simplify: if simplify:
try: try:
check_requirements(('onnx-simplifier',)) cuda = torch.cuda.is_available()
check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
import onnxsim import onnxsim
LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
model_onnx, check = onnxsim.simplify(model_onnx, model_onnx, check = onnxsim.simplify(model_onnx)
dynamic_input_shape=dynamic,
input_shapes={'images': list(im.shape)} if dynamic else None)
assert check, 'assert check failed' assert check, 'assert check failed'
onnx.save(model_onnx, f) onnx.save(model_onnx, f)
except Exception as e: except Exception as e:
@ -493,7 +492,7 @@ def run(
imgsz *= 2 if len(imgsz) == 1 else 1 # expand imgsz *= 2 if len(imgsz) == 1 else 1 # expand
assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}' assert nc == len(names), f'Model class count {nc} != len(names) {len(names)}'
if optimize: if optimize:
assert device.type != 'cuda', '--optimize not compatible with cuda devices, i.e. use --device cpu' assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu'
# Input # Input
gs = int(max(model.stride)) # grid size (max stride) gs = int(max(model.stride)) # grid size (max stride)

2
requirements.txt
View File

@ -25,7 +25,7 @@ seaborn>=0.11.0
# Export -------------------------------------- # Export --------------------------------------
# coremltools>=4.1 # CoreML export # coremltools>=4.1 # CoreML export
# onnx>=1.9.0 # ONNX export # onnx>=1.9.0 # ONNX export
# onnx-simplifier>=0.3.6 # ONNX simplifier # onnx-simplifier>=0.4.1 # ONNX simplifier
# nvidia-pyindex # TensorRT export # nvidia-pyindex # TensorRT export
# nvidia-tensorrt # TensorRT export # nvidia-tensorrt # TensorRT export
# scikit-learn==0.19.2 # CoreML quantization # scikit-learn==0.19.2 # CoreML quantization